Skip to content

Commit 53ec325

Browse files
KiteniteKitenite
authored
refactor: clean up message context management (#2952)
1 parent 407fd26 commit 53ec325

File tree

25 files changed

+2905
-187
lines changed

25 files changed

+2905
-187
lines changed

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

Lines changed: 14 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,12 @@
1+
import { getContextClass, getContextLabel } from '@onlook/ai';
12
import { DefaultSettings } from '@onlook/constants';
23
import { MessageContextType, type MessageContext } from '@onlook/models/chat';
3-
import { Icons } from '@onlook/ui/icons';
44
import { getTruncatedFileName } from '@onlook/ui/utils';
5-
import { assertNever } from '@onlook/utility';
6-
import React from 'react';
75
import { NodeIcon } from '../../../left-panel/layers-tab/tree/node-icon';
86

97
export function getTruncatedName(context: MessageContext) {
10-
let name = context.displayName;
8+
let name = getContextLabel(context);
9+
1110
if (context.type === MessageContextType.FILE || context.type === MessageContextType.IMAGE) {
1211
name = getTruncatedFileName(name);
1312
}
@@ -18,33 +17,19 @@ export function getTruncatedName(context: MessageContext) {
1817
}
1918

2019
export function getContextIcon(context: MessageContext) {
21-
let icon: React.ComponentType | React.ReactElement | null = null;
22-
switch (context.type) {
23-
case MessageContextType.FILE:
24-
icon = Icons.File;
25-
break;
26-
case MessageContextType.IMAGE:
27-
icon = Icons.Image;
28-
break;
29-
case MessageContextType.ERROR:
30-
icon = Icons.InfoCircled;
31-
break;
32-
case MessageContextType.HIGHLIGHT:
33-
return (
34-
<NodeIcon tagName={context.displayName} iconClass="w-3 h-3 ml-1 mr-2 flex-none" />
35-
);
36-
case MessageContextType.BRANCH:
37-
icon = Icons.Branch;
38-
break;
39-
case MessageContextType.AGENT_RULE:
40-
icon = Icons.Cube;
41-
break;
42-
default:
43-
assertNever(context);
20+
// Special case for highlight context which uses a custom component
21+
if (context.type === MessageContextType.HIGHLIGHT) {
22+
return (
23+
<NodeIcon tagName={context.displayName} iconClass="w-3 h-3 ml-1 mr-2 flex-none" />
24+
);
4425
}
45-
if (icon) {
46-
return React.createElement(icon);
26+
27+
const contextClass = getContextClass(context.type);
28+
if (contextClass?.icon) {
29+
const IconComponent = contextClass.icon;
30+
return <IconComponent />;
4731
}
32+
return null;
4833
}
4934

5035
export function validateImageLimit(

‎packages/ai/package.json‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
{
22
"name": "@onlook/ai",
33
"description": "A AI library for Onlook",
4-
"main": "./dist/index.js",
54
"type": "module",
65
"module": "src/index.ts",
76
"types": "src/index.ts",

‎packages/ai/src/contexts/classes/agent-rule-context.ts‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import { MessageContextType, type AgentRuleMessageContext } from '@onlook/models';
2+
import { Icons } from '@onlook/ui/icons';
3+
import { wrapXml } from '../../prompt/helpers';
4+
import { BaseContext } from '../models/base';
5+
6+
export class AgentRuleContext extends BaseContext {
7+
static readonly contextType = MessageContextType.AGENT_RULE;
8+
static readonly displayName = 'Agent Rule';
9+
static readonly icon = Icons.Cube;
10+
11+
private static readonly agentRulesContextPrefix = `These are user provided rules for the project`;
12+
13+
static getPrompt(context: AgentRuleMessageContext): string {
14+
return `${context.path}\n${context.content}`;
15+
}
16+
17+
static getLabel(context: AgentRuleMessageContext): string {
18+
return context.displayName || context.path;
19+
}
20+
21+
/**
22+
* Generate multiple agent rules content
23+
*/
24+
static getAgentRulesContent(agentRules: AgentRuleMessageContext[]): string {
25+
let content = `${AgentRuleContext.agentRulesContextPrefix}\n`;
26+
const rulePrompts = agentRules.map(agentRule => AgentRuleContext.getPrompt(agentRule));
27+
content += rulePrompts.join('\n');
28+
return wrapXml('agent-rules', content);
29+
}
30+
}

‎packages/ai/src/contexts/classes/branch-context.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import { MessageContextType, type BranchMessageContext } from '@onlook/models';
2+
import { Icons } from '@onlook/ui/icons';
3+
import { wrapXml } from '../../prompt/helpers';
4+
import { BaseContext } from '../models/base';
5+
6+
export class BranchContext extends BaseContext {
7+
static readonly contextType = MessageContextType.BRANCH;
8+
static readonly displayName = 'Branch';
9+
static readonly icon = Icons.Branch;
10+
11+
static getPrompt(context: BranchMessageContext): string {
12+
return `Branch: ${context.branch.name} (${context.branch.id})\nDescription: ${context.content}`;
13+
}
14+
15+
static getLabel(context: BranchMessageContext): string {
16+
return context.displayName || context.branch.name;
17+
}
18+
19+
/**
20+
* Generate multiple branches content
21+
*/
22+
static getBranchesContent(branches: BranchMessageContext[]): string {
23+
let prompt = `I'm working on the following branches: \n`;
24+
prompt += branches.map((b) => b.branch.id).join(', ');
25+
prompt = wrapXml('branches', prompt);
26+
return prompt;
27+
}
28+
}

‎packages/ai/src/contexts/classes/error-context.ts‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import { MessageContextType, type ErrorMessageContext } from '@onlook/models';
2+
import { Icons } from '@onlook/ui/icons';
3+
import { wrapXml } from '../../prompt/helpers';
4+
import { BaseContext } from '../models/base';
5+
6+
export class ErrorContext extends BaseContext {
7+
static readonly contextType = MessageContextType.ERROR;
8+
static readonly displayName = 'Error';
9+
static readonly icon = Icons.InfoCircled;
10+
11+
private static readonly errorsContentPrefix = `You are helping debug a Next.js React app, likely being set up for the first time. Common issues:
12+
- Missing dependencies ("command not found" errors) → Suggest "bun install" to install the dependencies for the first time (this project uses Bun, not npm)
13+
- Missing closing tags in JSX/TSX files. Make sure all the tags are closed.
14+
15+
The errors can be from terminal or browser and might have the same root cause. Analyze all the messages before suggesting solutions. If there is no solution, don't suggest a fix.
16+
If the same error is being reported multiple times, the previous fix did not work. Try a different approach.
17+
18+
IMPORTANT: This project uses Bun as the package manager. Always use Bun commands:
19+
- Use "bun install" instead of "npm install"
20+
- Use "bun add" instead of "npm install <package>"
21+
- Use "bun run" instead of "npm run"
22+
- Use "bunx" instead of "npx"
23+
24+
NEVER SUGGEST THE "bun run dev" command. Assume the user is already running the app.`;
25+
26+
static getPrompt(context: ErrorMessageContext): string {
27+
const branchDisplay = ErrorContext.getBranchContent(context.branchId);
28+
const errorDisplay = wrapXml('error', context.content);
29+
return `${branchDisplay}\n${errorDisplay}\n`;
30+
}
31+
32+
static getLabel(context: ErrorMessageContext): string {
33+
return context.displayName || 'Error';
34+
}
35+
36+
/**
37+
* Generate multiple errors content
38+
*/
39+
static getErrorsContent(errors: ErrorMessageContext[]): string {
40+
if (errors.length === 0) {
41+
return '';
42+
}
43+
let prompt = `${ErrorContext.errorsContentPrefix}\n`;
44+
for (const error of errors) {
45+
prompt += ErrorContext.getPrompt(error);
46+
}
47+
48+
prompt = wrapXml('errors', prompt);
49+
return prompt;
50+
}
51+
52+
private static getBranchContent(id: string): string {
53+
return wrapXml('branch', `id: "${id}"`);
54+
}
55+
}

‎packages/ai/src/contexts/classes/file-context.ts‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import { MessageContextType, type FileMessageContext, type HighlightMessageContext } from '@onlook/models';
2+
import { Icons } from '@onlook/ui/icons';
3+
import { CODE_FENCE } from '../../prompt/constants';
4+
import { wrapXml } from '../../prompt/helpers';
5+
import { BaseContext } from '../models/base';
6+
7+
export class FileContext extends BaseContext {
8+
static readonly contextType = MessageContextType.FILE;
9+
static readonly displayName = 'File';
10+
static readonly icon = Icons.File;
11+
12+
private static readonly filesContentPrefix = `I have added these files to the chat so you can go ahead and edit them`;
13+
private static readonly truncatedFilesContentPrefix = `This context originally included the content of files listed below and has been truncated to save space.
14+
If relevant, feel free to retrieve their content.`;
15+
16+
static getPrompt(context: FileMessageContext): string {
17+
const pathDisplay = wrapXml('path', context.path);
18+
const branchDisplay = wrapXml('branch', `id: "${context.branchId}"`);
19+
let prompt = `${pathDisplay}\n${branchDisplay}\n`;
20+
prompt += `${CODE_FENCE.start}${FileContext.getLanguageFromFilePath(context.path)}\n`;
21+
prompt += context.content;
22+
prompt += `\n${CODE_FENCE.end}\n`;
23+
return prompt;
24+
}
25+
26+
static getLabel(context: FileMessageContext): string {
27+
return context.path.split('/').pop() || 'File';
28+
}
29+
30+
/**
31+
* Generate multiple files content with highlights
32+
*/
33+
static getFilesContent(files: FileMessageContext[], highlights: HighlightMessageContext[]): string {
34+
if (files.length === 0) {
35+
return '';
36+
}
37+
let prompt = '';
38+
prompt += `${FileContext.filesContentPrefix}\n`;
39+
let index = 1;
40+
for (const file of files) {
41+
let filePrompt = FileContext.getPrompt(file);
42+
// Add highlights for this file
43+
const highlightContent = FileContext.getHighlightsForFile(file.path, highlights, file.branchId);
44+
if (highlightContent) {
45+
filePrompt += highlightContent;
46+
}
47+
48+
filePrompt = wrapXml(files.length > 1 ? `file-${index}` : 'file', filePrompt);
49+
prompt += filePrompt;
50+
index++;
51+
}
52+
53+
return prompt;
54+
}
55+
56+
/**
57+
* Generate truncated files content
58+
*/
59+
static getTruncatedFilesContent(files: FileMessageContext[]): string {
60+
if (files.length === 0) {
61+
return '';
62+
}
63+
let prompt = '';
64+
prompt += `${FileContext.truncatedFilesContentPrefix}\n`;
65+
let index = 1;
66+
for (const file of files) {
67+
const branchDisplay = FileContext.getBranchContent(file.branchId);
68+
const pathDisplay = wrapXml('path', file.path);
69+
let filePrompt = `${pathDisplay}\n${branchDisplay}\n`;
70+
filePrompt = wrapXml(files.length > 1 ? `file-${index}` : 'file', filePrompt);
71+
prompt += filePrompt;
72+
index++;
73+
}
74+
75+
return prompt;
76+
}
77+
78+
private static getBranchContent(id: string): string {
79+
return wrapXml('branch', `id: "${id}"`);
80+
}
81+
82+
private static getHighlightsForFile(filePath: string, highlights: HighlightMessageContext[], branchId: string): string {
83+
// Import HighlightContext dynamically to avoid circular imports
84+
const { HighlightContext } = require('./highlight-context');
85+
return HighlightContext.getHighlightsContent(filePath, highlights, branchId);
86+
}
87+
88+
private static getLanguageFromFilePath(filePath: string): string {
89+
return filePath.split('.').pop() ?? '';
90+
}
91+
}

‎packages/ai/src/contexts/classes/highlight-context.ts‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import { MessageContextType, type HighlightMessageContext } from '@onlook/models';
2+
import { Icons } from '@onlook/ui/icons';
3+
import { CODE_FENCE } from '../../prompt/constants';
4+
import { wrapXml } from '../../prompt/helpers';
5+
import { BaseContext } from '../models/base';
6+
7+
export class HighlightContext extends BaseContext {
8+
static readonly contextType = MessageContextType.HIGHLIGHT;
9+
static readonly displayName = 'Code Selection';
10+
static readonly icon = Icons.CursorArrow;
11+
12+
private static readonly highlightPrefix = '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.';
13+
14+
static getPrompt(context: HighlightMessageContext): string {
15+
const branchDisplay = HighlightContext.getBranchContent(context.branchId);
16+
const pathDisplay = wrapXml('path', `${context.path}#L${context.start}:L${context.end}`);
17+
let prompt = `${pathDisplay}\n${branchDisplay}\n`;
18+
prompt += `${CODE_FENCE.start}\n`;
19+
prompt += context.content;
20+
prompt += `\n${CODE_FENCE.end}\n`;
21+
return prompt;
22+
}
23+
24+
static getLabel(context: HighlightMessageContext): string {
25+
return context.displayName || context.path.split('/').pop() || 'Code Selection';
26+
}
27+
28+
/**
29+
* Generate multiple highlights content for a file path
30+
*/
31+
static getHighlightsContent(filePath: string, highlights: HighlightMessageContext[], branchId: string): string {
32+
const fileHighlights = highlights.filter((h) => h.path === filePath && h.branchId === branchId);
33+
if (fileHighlights.length === 0) {
34+
return '';
35+
}
36+
let prompt = `${HighlightContext.highlightPrefix}\n`;
37+
let index = 1;
38+
for (const highlight of fileHighlights) {
39+
let highlightPrompt = HighlightContext.getPrompt(highlight);
40+
highlightPrompt = wrapXml(
41+
fileHighlights.length > 1 ? `highlight-${index}` : 'highlight',
42+
highlightPrompt,
43+
);
44+
prompt += highlightPrompt;
45+
index++;
46+
}
47+
return prompt;
48+
}
49+
50+
private static getBranchContent(id: string): string {
51+
return wrapXml('branch', `id: "${id}"`);
52+
}
53+
}

‎packages/ai/src/contexts/classes/image-context.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import { MessageContextType, type ImageMessageContext } from '@onlook/models';
2+
import { Icons } from '@onlook/ui/icons';
3+
import { BaseContext } from '../models/base';
4+
5+
export class ImageContext extends BaseContext {
6+
static readonly contextType = MessageContextType.IMAGE;
7+
static readonly displayName = 'Image';
8+
static readonly icon = Icons.Image;
9+
10+
static getPrompt(context: ImageMessageContext): string {
11+
// Images don't generate text prompts - they're handled as file attachments
12+
return `[Image: ${context.mimeType}]`;
13+
}
14+
15+
static getLabel(context: ImageMessageContext): string {
16+
return context.displayName || 'Image';
17+
}
18+
19+
/**
20+
* Convert image contexts to file UI parts for AI SDK
21+
*/
22+
static toFileUIParts(images: ImageMessageContext[]) {
23+
return images.map((i) => ({
24+
type: 'file' as const,
25+
mediaType: i.mimeType,
26+
url: i.content,
27+
}));
28+
}
29+
}

0 commit comments

Comments
 (0)