Skip to content

Commit 28cde75

Browse files
KiteniteKitenite
authored
fix: snapshot type conflict and mastra can't view image (#2656)
* Revert "fix: typecheck snapshot (#2655)" This reverts commit 782fd74. * properly deprecate snapshot * revert mastra for image viewing * refactor helpers
1 parent 782fd74 commit 28cde75

File tree

28 files changed

+221
-185
lines changed

28 files changed

+221
-185
lines changed

‎apps/web/client/src/app/api/chat/helperts/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export * from './stream';
2+
export * from './usage';

‎apps/web/client/src/app/api/chat/helperts/stream.ts‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
import { ASK_TOOL_SET, BUILD_TOOL_SET, getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, initModel } from '@onlook/ai';
2+
import { ChatType, LLMProvider, OPENROUTER_MODELS, type ModelConfig } from '@onlook/models';
3+
import { generateObject, NoSuchToolError, type ToolCall, type ToolSet } from 'ai';
4+
5+
export async function getModelFromType(chatType: ChatType) {
6+
let model: ModelConfig;
7+
switch (chatType) {
8+
case ChatType.CREATE:
9+
case ChatType.FIX:
10+
model = await initModel({
11+
provider: LLMProvider.OPENROUTER,
12+
model: OPENROUTER_MODELS.OPEN_AI_GPT_5,
13+
});
14+
break;
15+
case ChatType.ASK:
16+
case ChatType.EDIT:
17+
default:
18+
model = await initModel({
19+
provider: LLMProvider.OPENROUTER,
20+
model: OPENROUTER_MODELS.CLAUDE_4_SONNET,
21+
});
22+
break;
23+
}
24+
return model;
25+
}
26+
27+
export async function getToolSetFromType(chatType: ChatType) {
28+
return chatType === ChatType.ASK ? ASK_TOOL_SET : BUILD_TOOL_SET;
29+
}
30+
31+
export async function getSystemPromptFromType(chatType: ChatType) {
32+
let systemPrompt: string;
33+
34+
switch (chatType) {
35+
case ChatType.CREATE:
36+
systemPrompt = getCreatePageSystemPrompt();
37+
break;
38+
case ChatType.ASK:
39+
systemPrompt = getAskModeSystemPrompt();
40+
break;
41+
case ChatType.EDIT:
42+
default:
43+
systemPrompt = getSystemPrompt();
44+
break;
45+
}
46+
return systemPrompt;
47+
}
48+
49+
50+
export const repairToolCall = async ({ toolCall, tools, error }: { toolCall: ToolCall<string, any>, tools: ToolSet, error: Error }) => {
51+
if (NoSuchToolError.isInstance(error)) {
52+
throw new Error(
53+
`Tool "${toolCall.toolName}" not found. Available tools: ${Object.keys(tools).join(', ')}`,
54+
);
55+
}
56+
const tool = tools[toolCall.toolName as keyof typeof tools];
57+
58+
console.warn(
59+
`Invalid parameter for tool ${toolCall.toolName} with args ${JSON.stringify(toolCall.args)}, attempting to fix`,
60+
);
61+
62+
const { model } = await initModel({
63+
provider: LLMProvider.OPENROUTER,
64+
model: OPENROUTER_MODELS.CLAUDE_4_SONNET,
65+
});
66+
67+
const { object: repairedArgs } = await generateObject({
68+
model,
69+
schema: tool?.parameters,
70+
prompt: [
71+
`The model tried to call the tool "${toolCall.toolName}"` +
72+
` with the following arguments:`,
73+
JSON.stringify(toolCall.args),
74+
`The tool accepts the following schema:`,
75+
JSON.stringify(tool?.parameters),
76+
'Please fix the arguments.',
77+
].join('\n'),
78+
});
79+
80+
return {
81+
...toolCall,
82+
args: JSON.stringify(repairedArgs),
83+
toolCallType: 'function' as const
84+
};
85+
}
86+
87+
export function errorHandler(error: unknown) {
88+
try {
89+
console.error('Error in chat', error);
90+
if (!error) {
91+
return 'unknown error';
92+
}
93+
94+
if (typeof error === 'string') {
95+
return error;
96+
}
97+
98+
if (error instanceof Error) {
99+
return error.message;
100+
}
101+
return JSON.stringify(error);
102+
} catch (error) {
103+
console.error('Error in errorHandler', error);
104+
return 'unknown error';
105+
}
106+
}

‎apps/web/client/src/app/api/chat/helpers.ts‎ renamed to ‎apps/web/client/src/app/api/chat/helperts/usage.ts‎

Lines changed: 0 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -42,43 +42,6 @@ export const getSupabaseUser = async (request: NextRequest) => {
4242
return user;
4343
}
4444

45-
export const repairToolCall = async ({ toolCall, tools, error }: { toolCall: ToolCall<string, any>, tools: ToolSet, error: Error }) => {
46-
if (NoSuchToolError.isInstance(error)) {
47-
throw new Error(
48-
`Tool "${toolCall.toolName}" not found. Available tools: ${Object.keys(tools).join(', ')}`,
49-
);
50-
}
51-
const tool = tools[toolCall.toolName as keyof typeof tools];
52-
53-
console.warn(
54-
`Invalid parameter for tool ${toolCall.toolName} with args ${JSON.stringify(toolCall.args)}, attempting to fix`,
55-
);
56-
57-
const { model } = await initModel({
58-
provider: LLMProvider.OPENROUTER,
59-
model: OPENROUTER_MODELS.CLAUDE_4_SONNET,
60-
});
61-
62-
const { object: repairedArgs } = await generateObject({
63-
model,
64-
schema: tool?.parameters,
65-
prompt: [
66-
`The model tried to call the tool "${toolCall.toolName}"` +
67-
` with the following arguments:`,
68-
JSON.stringify(toolCall.args),
69-
`The tool accepts the following schema:`,
70-
JSON.stringify(tool?.parameters),
71-
'Please fix the arguments.',
72-
].join('\n'),
73-
});
74-
75-
return {
76-
...toolCall,
77-
args: JSON.stringify(repairedArgs),
78-
toolCallType: 'function' as const
79-
};
80-
}
81-
8245
export const incrementUsage = async (req: NextRequest): Promise<{
8346
usageRecordId: string | undefined,
8447
rateLimitId: string | undefined,

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

Lines changed: 18 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
1-
import { mastra } from '@/mastra';
2-
import { CHAT_TYPE_KEY, ONLOOK_AGENT_KEY, type OnlookAgentRuntimeContext } from '@/mastra/agents';
31
import { trackEvent } from '@/utils/analytics/server';
4-
import { RuntimeContext } from '@mastra/core/runtime-context';
52
import { convertToStreamMessages } from '@onlook/ai';
63
import { ChatType, type ChatMessage } from '@onlook/models';
4+
import { streamText } from 'ai';
75
import { type NextRequest } from 'next/server';
8-
import { checkMessageLimit, decrementUsage, getSupabaseUser, incrementUsage, repairToolCall } from './helpers';
6+
import { checkMessageLimit, decrementUsage, errorHandler, getModelFromType, getSupabaseUser, getSystemPromptFromType, getToolSetFromType, incrementUsage, repairToolCall } from './helperts';
97

108
export async function POST(req: NextRequest) {
119
try {
@@ -54,9 +52,6 @@ export async function POST(req: NextRequest) {
5452

5553
export const streamResponse = async (req: NextRequest) => {
5654
const { messages, maxSteps, chatType }: { messages: ChatMessage[], maxSteps: number, chatType: ChatType } = await req.json();
57-
const agent = mastra.getAgent(ONLOOK_AGENT_KEY);
58-
const runtimeContext = new RuntimeContext<OnlookAgentRuntimeContext>()
59-
runtimeContext.set(CHAT_TYPE_KEY, chatType);
6055

6156
// Updating the usage record and rate limit is done here to avoid
6257
// abuse in the case where a single user sends many concurrent requests.
@@ -68,14 +63,24 @@ export const streamResponse = async (req: NextRequest) => {
6863
if (chatType === ChatType.EDIT) {
6964
usageRecord = await incrementUsage(req);
7065
}
71-
const result = await agent.stream(convertToStreamMessages(messages), {
72-
headers: {
73-
'HTTP-Referer': 'https://onlook.com',
74-
'X-Title': 'Onlook',
75-
},
66+
67+
const { model, providerOptions, headers } = await getModelFromType(chatType);
68+
const systemPrompt = await getSystemPromptFromType(chatType);
69+
const tools = await getToolSetFromType(chatType);
70+
const result = streamText({
71+
model,
72+
headers,
73+
tools,
7674
maxSteps,
77-
runtimeContext,
7875
toolCallStreaming: true,
76+
messages: [
77+
{
78+
role: 'system',
79+
content: systemPrompt,
80+
providerOptions,
81+
},
82+
...convertToStreamMessages(messages),
83+
],
7984
experimental_repairToolCall: repairToolCall,
8085
onError: async (error) => {
8186
console.error('Error in chat', error);
@@ -90,24 +95,3 @@ export const streamResponse = async (req: NextRequest) => {
9095
}
9196
);
9297
}
93-
94-
export function errorHandler(error: unknown) {
95-
try {
96-
console.error('Error in chat', error);
97-
if (!error) {
98-
return 'unknown error';
99-
}
100-
101-
if (typeof error === 'string') {
102-
return error;
103-
}
104-
105-
if (error instanceof Error) {
106-
return error.message;
107-
}
108-
return JSON.stringify(error);
109-
} catch (error) {
110-
console.error('Error in errorHandler', error);
111-
return 'unknown error';
112-
}
113-
}

‎apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/user-message.tsx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { useChatContext } from '@/app/project/[id]/_hooks/use-chat';
22
import { useEditorEngine } from '@/components/store/editor';
3-
import { ChatType, MessageSnapshotType, type UserChatMessage } from '@onlook/models';
3+
import { ChatType, MessageCheckpointType, type UserChatMessage } from '@onlook/models';
44
import { Button } from '@onlook/ui/button';
55
import { Icons } from '@onlook/ui/icons';
66
import { toast } from '@onlook/ui/sonner';
@@ -35,7 +35,7 @@ export const UserMessage = ({ message }: UserMessageProps) => {
3535
const [isRestoring, setIsRestoring] = useState(false);
3636

3737
const textareaRef = useRef<HTMLTextAreaElement>(null);
38-
const commitOid = message.content.metadata?.snapshots?.find((s) => s.type === MessageSnapshotType.GIT)?.oid;
38+
const commitOid = message.content.metadata.checkpoints.find((s) => s.type === MessageCheckpointType.GIT)?.oid;
3939

4040
useEffect(() => {
4141
if (isEditing && textareaRef.current) {
@@ -210,7 +210,7 @@ export const UserMessage = ({ message }: UserMessageProps) => {
210210
<div className="h-6 relative">
211211
<div className="absolute top-1 left-0 right-0 flex flex-row justify-start items-center w-full overflow-auto pr-16">
212212
<div className="flex flex-row gap-3 text-micro text-foreground-secondary">
213-
{message.content.metadata?.context?.map((context) => (
213+
{message.content.metadata.context.map((context) => (
214214
<SentContextPill key={nanoid()} context={context} />
215215
))}
216216
</div>

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { type ChatMessageContext } from '@onlook/models/chat';
1+
import { type MessageContext } from '@onlook/models/chat';
22
import { Icons } from '@onlook/ui/icons';
33
import { motion } from 'motion/react';
44
import React from 'react';
@@ -7,7 +7,7 @@ import { getContextIcon, getTruncatedName } from './helpers';
77
export const DraftContextPill = React.forwardRef<
88
HTMLDivElement,
99
{
10-
context: ChatMessageContext;
10+
context: MessageContext;
1111
onRemove: () => void;
1212
}
1313
>(({ context, onRemove }, ref) => {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { type ChatMessageContext, MessageContextType } from '@onlook/models/chat';
1+
import { type MessageContext, MessageContextType } from '@onlook/models/chat';
22
import { Icons } from '@onlook/ui/icons';
33
import { motion } from 'motion/react';
44
import React from 'react';
@@ -7,7 +7,7 @@ import { getTruncatedName } from './helpers';
77
export const DraftImagePill = React.forwardRef<
88
HTMLDivElement,
99
{
10-
context: ChatMessageContext;
10+
context: MessageContext;
1111
onRemove: () => void;
1212
}
1313
>(({ context, onRemove }, ref) => {

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
import { MessageContextType, type ChatMessageContext } from '@onlook/models/chat';
1+
import { MessageContextType, type MessageContext } from '@onlook/models/chat';
22
import { Icons } from '@onlook/ui/icons';
33
import { getTruncatedFileName } from '@onlook/ui/utils';
44
import { assertNever } from '@onlook/utility';
55
import React from 'react';
66
import { NodeIcon } from '../../../left-panel/layers-tab/tree/node-icon';
77

8-
export function getTruncatedName(context: ChatMessageContext) {
8+
export function getTruncatedName(context: MessageContext) {
99
let name = context.displayName;
1010
if (context.type === MessageContextType.FILE || context.type === MessageContextType.IMAGE) {
1111
name = getTruncatedFileName(name);
@@ -16,7 +16,7 @@ export function getTruncatedName(context: ChatMessageContext) {
1616
return name.length > 20 ? `${name.slice(0, 20)}...` : name;
1717
}
1818

19-
export function getContextIcon(context: ChatMessageContext) {
19+
export function getContextIcon(context: MessageContext) {
2020
let icon: React.ComponentType | React.ReactElement | null = null;
2121
switch (context.type) {
2222
case MessageContextType.FILE:

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { useEditorEngine } from '@/components/store/editor';
2-
import type { ChatMessageContext } from '@onlook/models/chat';
2+
import type { MessageContext } from '@onlook/models/chat';
33
import { MessageContextType } from '@onlook/models/chat';
44
import { cn } from '@onlook/ui/utils';
55
import { observer } from 'mobx-react-lite';
@@ -10,7 +10,7 @@ import { DraftImagePill } from './draft-image-pill';
1010
export const InputContextPills = observer(() => {
1111
const editorEngine = useEditorEngine();
1212

13-
const handleRemoveContext = (contextToRemove: ChatMessageContext) => {
13+
const handleRemoveContext = (contextToRemove: MessageContext) => {
1414
const newContext = [...editorEngine.chat.context.context].filter(
1515
(context) => context !== contextToRemove,
1616
);
@@ -27,7 +27,7 @@ export const InputContextPills = observer(() => {
2727
>
2828
<AnimatePresence mode="popLayout">
2929
{editorEngine.chat.context.context.map(
30-
(context: ChatMessageContext, index: number) => {
30+
(context: MessageContext, index: number) => {
3131
if (context.type === MessageContextType.IMAGE) {
3232
return (
3333
<DraftImagePill

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
import { type ChatMessageContext } from '@onlook/models/chat';
1+
import { type MessageContext } from '@onlook/models/chat';
22
import { getContextIcon, getTruncatedName } from './helpers';
33

4-
export function SentContextPill({ context }: { context: ChatMessageContext }) {
4+
export function SentContextPill({ context }: { context: MessageContext }) {
55
return (
66
<span
77
className="flex flex-row gap-0.5 text-xs items-center select-none"

0 commit comments

Comments
 (0)