diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts index 2149e2b..e79bf4b 100644 --- a/app/(chat)/api/chat/route.ts +++ b/app/(chat)/api/chat/route.ts @@ -13,9 +13,7 @@ import { type ResumableStreamContext, } from "resumable-stream"; import { auth, type UserType } from "@/app/(auth)/auth"; -import type { VisibilityType } from "@/components/visibility-selector"; import { entitlementsByUserType } from "@/lib/ai/entitlements"; -import type { ChatModel } from "@/lib/ai/models"; import { type RequestHints, systemPrompt } from "@/lib/ai/prompts"; import { getLanguageModel } from "@/lib/ai/providers"; import { createDocument } from "@/lib/ai/tools/create-document"; @@ -32,6 +30,7 @@ import { saveChat, saveMessages, updateChatTitleById, + updateMessage, } from "@/lib/db/queries"; import type { DBMessage } from "@/lib/db/schema"; import { ChatSDKError } from "@/lib/errors"; @@ -75,17 +74,8 @@ export async function POST(request: Request) { } try { - const { - id, - message, - selectedChatModel, - selectedVisibilityType, - }: { - id: string; - message: ChatMessage; - selectedChatModel: ChatModel["id"]; - selectedVisibilityType: VisibilityType; - } = requestBody; + const { id, message, messages, selectedChatModel, selectedVisibilityType } = + requestBody; const session = await auth(); @@ -104,6 +94,9 @@ export async function POST(request: Request) { return new ChatSDKError("rate_limit:chat").toResponse(); } + // Check if this is a tool approval flow (all messages sent) + const isToolApprovalFlow = Boolean(messages); + const chat = await getChatById({ id }); let messagesFromDb: DBMessage[] = []; let titlePromise: Promise | null = null; @@ -112,9 +105,11 @@ export async function POST(request: Request) { if (chat.userId !== session.user.id) { return new ChatSDKError("forbidden:chat").toResponse(); } - // Only fetch messages if chat already exists - messagesFromDb = await getMessagesByChatId({ id }); - } else { + // Only fetch messages if chat already exists and not tool approval + if (!isToolApprovalFlow) { + messagesFromDb = await getMessagesByChatId({ id }); + } + } else if (message?.role === "user") { // Save chat immediately with placeholder title await saveChat({ id, @@ -127,7 +122,10 @@ export async function POST(request: Request) { titlePromise = generateTitleFromUserMessage({ message }); } - const uiMessages = [...convertToUIMessages(messagesFromDb), message]; + // Use all messages for tool approval, otherwise DB messages + new message + const uiMessages = isToolApprovalFlow + ? (messages as ChatMessage[]) + : [...convertToUIMessages(messagesFromDb), message as ChatMessage]; const { longitude, latitude, city, country } = geolocation(request); @@ -138,24 +136,29 @@ export async function POST(request: Request) { country, }; - await saveMessages({ - messages: [ - { - chatId: id, - id: message.id, - role: "user", - parts: message.parts, - attachments: [], - createdAt: new Date(), - }, - ], - }); + // Only save user messages to the database (not tool approval responses) + if (message?.role === "user") { + await saveMessages({ + messages: [ + { + chatId: id, + id: message.id, + role: "user", + parts: message.parts, + attachments: [], + createdAt: new Date(), + }, + ], + }); + } const streamId = generateUUID(); await createStreamId({ streamId, chatId: id }); const stream = createUIMessageStream({ - execute: ({ writer: dataStream }) => { + // Pass original messages for tool approval continuation + originalMessages: isToolApprovalFlow ? uiMessages : undefined, + execute: async ({ writer: dataStream }) => { // Handle title generation in parallel if (titlePromise) { titlePromise.then((title) => { @@ -171,7 +174,7 @@ export async function POST(request: Request) { const result = streamText({ model: getLanguageModel(selectedChatModel), system: systemPrompt({ selectedChatModel, requestHints }), - messages: convertToModelMessages(uiMessages), + messages: await convertToModelMessages(uiMessages), stopWhen: stepCountIs(5), experimental_activeTools: isReasoningModel ? [] @@ -215,32 +218,67 @@ export async function POST(request: Request) { ); }, generateId: generateUUID, - onFinish: async ({ messages }) => { - await saveMessages({ - messages: messages.map((currentMessage) => ({ - id: currentMessage.id, - role: currentMessage.role, - parts: currentMessage.parts, - createdAt: new Date(), - attachments: [], - chatId: id, - })), - }); + onFinish: async ({ messages: finishedMessages }) => { + if (isToolApprovalFlow) { + // For tool approval, update existing messages (tool state changed) and save new ones + for (const finishedMsg of finishedMessages) { + const existingMsg = uiMessages.find((m) => m.id === finishedMsg.id); + if (existingMsg) { + // Update existing message with new parts (tool state changed) + await updateMessage({ + id: finishedMsg.id, + parts: finishedMsg.parts, + }); + } else { + // Save new message + await saveMessages({ + messages: [ + { + id: finishedMsg.id, + role: finishedMsg.role, + parts: finishedMsg.parts, + createdAt: new Date(), + attachments: [], + chatId: id, + }, + ], + }); + } + } + } else if (finishedMessages.length > 0) { + // Normal flow - save all finished messages + await saveMessages({ + messages: finishedMessages.map((currentMessage) => ({ + id: currentMessage.id, + role: currentMessage.role, + parts: currentMessage.parts, + createdAt: new Date(), + attachments: [], + chatId: id, + })), + }); + } }, onError: () => { return "Oops, an error occurred!"; }, }); - // const streamContext = getStreamContext(); + const streamContext = getStreamContext(); - // if (streamContext) { - // return new Response( - // await streamContext.resumableStream(streamId, () => - // stream.pipeThrough(new JsonToSseTransformStream()) - // ) - // ); - // } + if (streamContext) { + try { + const resumableStream = await streamContext.resumableStream( + streamId, + () => stream.pipeThrough(new JsonToSseTransformStream()) + ); + if (resumableStream) { + return new Response(resumableStream); + } + } catch (error) { + console.error("Failed to create resumable stream:", error); + } + } return new Response(stream.pipeThrough(new JsonToSseTransformStream())); } catch (error) { diff --git a/app/(chat)/api/chat/schema.ts b/app/(chat)/api/chat/schema.ts index 2bb383e..60a708a 100644 --- a/app/(chat)/api/chat/schema.ts +++ b/app/(chat)/api/chat/schema.ts @@ -14,13 +14,24 @@ const filePartSchema = z.object({ const partSchema = z.union([textPartSchema, filePartSchema]); +const userMessageSchema = z.object({ + id: z.string().uuid(), + role: z.enum(["user"]), + parts: z.array(partSchema), +}); + +// For tool approval flows, we accept all messages (more permissive schema) +const messageSchema = z.object({ + id: z.string(), + role: z.string(), + parts: z.array(z.any()), +}); + export const postRequestBodySchema = z.object({ id: z.string().uuid(), - message: z.object({ - id: z.string().uuid(), - role: z.enum(["user"]), - parts: z.array(partSchema), - }), + // Either a single new message or all messages (for tool approvals) + message: userMessageSchema.optional(), + messages: z.array(messageSchema).optional(), selectedChatModel: z.string(), selectedVisibilityType: z.enum(["public", "private"]), }); diff --git a/components/ai-elements/confirmation.tsx b/components/ai-elements/confirmation.tsx index fb59672..ee00b25 100644 --- a/components/ai-elements/confirmation.tsx +++ b/components/ai-elements/confirmation.tsx @@ -97,7 +97,6 @@ export const ConfirmationRequest = ({ children }: ConfirmationRequestProps) => { const { state } = useConfirmation(); // Only show when approval is requested - // @ts-expect-error state only available in AI SDK v6 if (state !== "approval-requested") { return null; } @@ -117,9 +116,7 @@ export const ConfirmationAccepted = ({ // Only show when approved and in response states if ( !approval?.approved || - // @ts-expect-error state only available in AI SDK v6 (state !== "approval-responded" && - // @ts-expect-error state only available in AI SDK v6 state !== "output-denied" && state !== "output-available") ) { @@ -141,9 +138,7 @@ export const ConfirmationRejected = ({ // Only show when rejected and in response states if ( approval?.approved !== false || - // @ts-expect-error state only available in AI SDK v6 (state !== "approval-responded" && - // @ts-expect-error state only available in AI SDK v6 state !== "output-denied" && state !== "output-available") ) { @@ -162,7 +157,6 @@ export const ConfirmationActions = ({ const { state } = useConfirmation(); // Only show when approval is requested - // @ts-expect-error state only available in AI SDK v6 if (state !== "approval-requested") { return null; } diff --git a/components/ai-elements/tool.tsx b/components/ai-elements/tool.tsx index 6032478..3df5d10 100644 --- a/components/ai-elements/tool.tsx +++ b/components/ai-elements/tool.tsx @@ -40,7 +40,6 @@ const getStatusBadge = (status: ToolUIPart["state"]) => { const labels: Record = { "input-streaming": "Pending", "input-available": "Running", - // @ts-expect-error state only available in AI SDK v6 "approval-requested": "Awaiting Approval", "approval-responded": "Responded", "output-available": "Completed", @@ -51,7 +50,6 @@ const getStatusBadge = (status: ToolUIPart["state"]) => { const icons: Record = { "input-streaming": , "input-available": , - // @ts-expect-error state only available in AI SDK v6 "approval-requested": , "approval-responded": , "output-available": , diff --git a/components/artifact-messages.tsx b/components/artifact-messages.tsx index 1c65230..3cfa972 100644 --- a/components/artifact-messages.tsx +++ b/components/artifact-messages.tsx @@ -9,6 +9,7 @@ import type { UIArtifact } from "./artifact"; import { PreviewMessage, ThinkingMessage } from "./message"; type ArtifactMessagesProps = { + addToolApprovalResponse: UseChatHelpers["addToolApprovalResponse"]; chatId: string; status: UseChatHelpers["status"]; votes: Vote[] | undefined; @@ -20,6 +21,7 @@ type ArtifactMessagesProps = { }; function PureArtifactMessages({ + addToolApprovalResponse, chatId, status, votes, @@ -45,6 +47,7 @@ function PureArtifactMessages({ > {messages.map((message, index) => ( - {status === "submitted" && } + {status === "submitted" && + !messages.some((msg) => + msg.parts?.some( + (part) => "state" in part && part.state === "approval-responded" + ) + ) && } ["addToolApprovalResponse"]; chatId: string; input: string; setInput: Dispatch>; @@ -320,6 +322,7 @@ function PureArtifact({
({ id, messages: initialMessages, experimental_throttle: 100, generateId: generateUUID, + // Auto-continue after tool approval (only for APPROVED tools) + // Denied tools don't need server continuation - state is saved on next user message + sendAutomaticallyWhen: ({ messages: currentMessages }) => { + const lastMessage = currentMessages.at(-1); + // Only continue if a tool was APPROVED (not denied) + const shouldContinue = + lastMessage?.parts?.some( + (part) => + "state" in part && + part.state === "approval-responded" && + "approval" in part && + (part.approval as { approved?: boolean })?.approved === true + ) ?? false; + return shouldContinue; + }, transport: new DefaultChatTransport({ api: "/api/chat", fetch: fetchWithErrorHandlers, prepareSendMessagesRequest(request) { + const lastMessage = request.messages.at(-1); + + // Check if this is a tool approval continuation: + // - Last message is NOT a user message (meaning no new user input) + // - OR any message has tool parts that were responded to (approved or denied) + const isToolApprovalContinuation = + lastMessage?.role !== "user" || + request.messages.some((msg) => + msg.parts?.some((part) => { + const state = (part as { state?: string }).state; + return ( + state === "approval-responded" || state === "output-denied" + ); + }) + ); + return { body: { id: request.id, - message: request.messages.at(-1), + // Send all messages for tool approval continuation, otherwise just the last user message + ...(isToolApprovalContinuation + ? { messages: request.messages } + : { message: lastMessage }), selectedChatModel: currentModelIdRef.current, selectedVisibilityType: visibilityType, ...request.body, @@ -170,6 +205,7 @@ export function Chat({ /> +
( -
+
diff --git a/components/elements/tool.tsx b/components/elements/tool.tsx index 9d77dcc..1d7751d 100644 --- a/components/elements/tool.tsx +++ b/components/elements/tool.tsx @@ -35,19 +35,25 @@ export type ToolHeaderProps = { }; const getStatusBadge = (status: ToolUIPart["state"]) => { - const labels = { + const labels: Record = { "input-streaming": "Pending", "input-available": "Running", + "approval-requested": "Pending", + "approval-responded": "Approved", "output-available": "Completed", "output-error": "Error", - } as const; + "output-denied": "Denied", + }; - const icons = { + const icons: Record = { "input-streaming": , "input-available": , + "approval-requested": , + "approval-responded": , "output-available": , "output-error": , - } as const; + "output-denied": , + }; return ( ["addToolApprovalResponse"]; chatId: string; message: ChatMessage; vote: Vote | undefined; @@ -76,9 +78,10 @@ const PurePreviewMessage = ({ ), "w-full": (message.role === "assistant" && - message.parts?.some( + (message.parts?.some( (p) => p.type === "text" && p.text?.trim() - )) || + ) || + message.parts?.some((p) => p.type.startsWith("tool-")))) || mode === "edit", "max-w-[calc(100%-2.5rem)] sm:max-w-[min(fit-content,80%)]": message.role === "user" && mode !== "edit", @@ -122,7 +125,7 @@ const PurePreviewMessage = ({
+ +
+ ); + } + + if (isDenied) { + return ( +
+ + + +
+ Weather lookup was denied. +
+
+
+
+ ); + } + + if (state === "approval-responded") { + return ( +
+ + + + + + +
+ ); + } return ( - - - - {state === "input-available" && ( - - )} - {state === "output-available" && ( - } - /> - )} - - +
+ + + + {(state === "input-available" || + state === "approval-requested") && ( + + )} + {state === "approval-requested" && approvalId && ( +
+ + +
+ )} +
+
+
); } @@ -285,22 +361,15 @@ const PurePreviewMessage = ({ export const PreviewMessage = memo( PurePreviewMessage, (prevProps, nextProps) => { - if (prevProps.isLoading !== nextProps.isLoading) { - return false; + if ( + prevProps.isLoading === nextProps.isLoading && + prevProps.message.id === nextProps.message.id && + prevProps.requiresScrollPadding === nextProps.requiresScrollPadding && + equal(prevProps.message.parts, nextProps.message.parts) && + equal(prevProps.vote, nextProps.vote) + ) { + return true; } - if (prevProps.message.id !== nextProps.message.id) { - return false; - } - if (prevProps.requiresScrollPadding !== nextProps.requiresScrollPadding) { - return false; - } - if (!equal(prevProps.message.parts, nextProps.message.parts)) { - return false; - } - if (!equal(prevProps.vote, nextProps.vote)) { - return false; - } - return false; } ); diff --git a/components/messages.tsx b/components/messages.tsx index d3653d5..729c430 100644 --- a/components/messages.tsx +++ b/components/messages.tsx @@ -10,6 +10,7 @@ import { Greeting } from "./greeting"; import { PreviewMessage, ThinkingMessage } from "./message"; type MessagesProps = { + addToolApprovalResponse: UseChatHelpers["addToolApprovalResponse"]; chatId: string; status: UseChatHelpers["status"]; votes: Vote[] | undefined; @@ -22,6 +23,7 @@ type MessagesProps = { }; function PureMessages({ + addToolApprovalResponse, chatId, status, votes, @@ -54,6 +56,7 @@ function PureMessages({ {messages.map((message, index) => ( ))} - {status === "submitted" && } + {status === "submitted" && + !messages.some((msg) => + msg.parts?.some( + (part) => "state" in part && part.state === "approval-responded" + ) + ) && }
= { * For users without an account */ guest: { - maxMessagesPerDay: 10, + maxMessagesPerDay: 20, }, /* diff --git a/lib/ai/models.mock.ts b/lib/ai/models.mock.ts index 7dc1c6a..31ee31a 100644 --- a/lib/ai/models.mock.ts +++ b/lib/ai/models.mock.ts @@ -6,6 +6,11 @@ const mockResponses: Record = { greeting: "Hello! How can I help you today?", }; +const mockUsage = { + inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 20, text: 20, reasoning: 0 }, +}; + function getResponseForPrompt(prompt: unknown): string { const promptStr = JSON.stringify(prompt).toLowerCase(); @@ -25,17 +30,14 @@ function getResponseForPrompt(prompt: unknown): string { const createMockModel = (): LanguageModel => { return { - specificationVersion: "v2", + specificationVersion: "v3", provider: "mock", modelId: "mock-model", defaultObjectGenerationMode: "tool", - supportedUrls: [], - supportsImageUrls: false, - supportsStructuredOutputs: false, + supportedUrls: {}, doGenerate: async ({ prompt }: { prompt: unknown }) => ({ - rawCall: { rawPrompt: null, rawSettings: {} }, finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 }, + usage: mockUsage, content: [{ type: "text", text: getResponseForPrompt(prompt) }], warnings: [], }), @@ -46,24 +48,26 @@ const createMockModel = (): LanguageModel => { return { stream: new ReadableStream({ async start(controller) { + controller.enqueue({ type: "text-start", id: "t1" }); for (const word of words) { controller.enqueue({ type: "text-delta", - textDelta: `${word} `, + id: "t1", + delta: `${word} `, }); await new Promise((resolve) => { setTimeout(resolve, 10); }); } + controller.enqueue({ type: "text-end", id: "t1" }); controller.enqueue({ type: "finish", finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20 }, + usage: mockUsage, }); controller.close(); }, }), - rawCall: { rawPrompt: null, rawSettings: {} }, }; }, } as unknown as LanguageModel; @@ -71,17 +75,14 @@ const createMockModel = (): LanguageModel => { const createMockReasoningModel = (): LanguageModel => { return { - specificationVersion: "v2", + specificationVersion: "v3", provider: "mock", modelId: "mock-reasoning-model", defaultObjectGenerationMode: "tool", - supportedUrls: [], - supportsImageUrls: false, - supportsStructuredOutputs: false, + supportedUrls: {}, doGenerate: async () => ({ - rawCall: { rawPrompt: null, rawSettings: {} }, finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 }, + usage: mockUsage, content: [{ type: "text", text: "This is a reasoned response." }], reasoning: [ { type: "text", text: "Let me think through this step by step..." }, @@ -91,62 +92,77 @@ const createMockReasoningModel = (): LanguageModel => { doStream: () => ({ stream: new ReadableStream({ async start(controller) { + controller.enqueue({ type: "reasoning-start", id: "r1" }); controller.enqueue({ - type: "reasoning", - textDelta: "Let me think through this step by step... ", + type: "reasoning-delta", + id: "r1", + delta: "Let me think through this step by step... ", }); + controller.enqueue({ type: "reasoning-end", id: "r1" }); await new Promise((resolve) => { setTimeout(resolve, 10); }); + controller.enqueue({ type: "text-start", id: "t1" }); controller.enqueue({ type: "text-delta", - textDelta: "This is a reasoned response.", + id: "t1", + delta: "This is a reasoned response.", }); + controller.enqueue({ type: "text-end", id: "t1" }); controller.enqueue({ type: "finish", finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20 }, + usage: mockUsage, }); controller.close(); }, }), - rawCall: { rawPrompt: null, rawSettings: {} }, }), } as unknown as LanguageModel; }; const createMockTitleModel = (): LanguageModel => { return { - specificationVersion: "v2", + specificationVersion: "v3", provider: "mock", modelId: "mock-title-model", defaultObjectGenerationMode: "tool", - supportedUrls: [], - supportsImageUrls: false, - supportsStructuredOutputs: false, + supportedUrls: {}, doGenerate: async () => ({ - rawCall: { rawPrompt: null, rawSettings: {} }, finishReason: "stop", - usage: { inputTokens: 5, outputTokens: 5, totalTokens: 10 }, + usage: { + inputTokens: { total: 5, noCache: 5, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 5, text: 5, reasoning: 0 }, + }, content: [{ type: "text", text: "Test Conversation" }], warnings: [], }), doStream: () => ({ stream: new ReadableStream({ start(controller) { + controller.enqueue({ type: "text-start", id: "t1" }); controller.enqueue({ type: "text-delta", - textDelta: "Test Conversation", + id: "t1", + delta: "Test Conversation", }); + controller.enqueue({ type: "text-end", id: "t1" }); controller.enqueue({ type: "finish", finishReason: "stop", - usage: { inputTokens: 5, outputTokens: 5 }, + usage: { + inputTokens: { + total: 5, + noCache: 5, + cacheRead: 0, + cacheWrite: 0, + }, + outputTokens: { total: 5, text: 5, reasoning: 0 }, + }, }); controller.close(); }, }), - rawCall: { rawPrompt: null, rawSettings: {} }, }), } as unknown as LanguageModel; }; diff --git a/lib/ai/models.test.ts b/lib/ai/models.test.ts index 1db7b41..66d7b81 100644 --- a/lib/ai/models.test.ts +++ b/lib/ai/models.test.ts @@ -1,12 +1,16 @@ import { simulateReadableStream } from "ai"; -import { MockLanguageModelV2 } from "ai/test"; +import { MockLanguageModelV3 } from "ai/test"; import { getResponseChunksByPrompt } from "@/tests/prompts/utils"; -export const chatModel = new MockLanguageModelV2({ +const mockUsage = { + inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 20, text: 20, reasoning: 0 }, +}; + +export const chatModel = new MockLanguageModelV3({ doGenerate: async () => ({ - rawCall: { rawPrompt: null, rawSettings: {} }, finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 }, + usage: mockUsage, content: [{ type: "text", text: "Hello, world!" }], warnings: [], }), @@ -16,15 +20,13 @@ export const chatModel = new MockLanguageModelV2({ initialDelayInMs: 1000, chunks: getResponseChunksByPrompt(prompt), }), - rawCall: { rawPrompt: null, rawSettings: {} }, }), }); -export const reasoningModel = new MockLanguageModelV2({ +export const reasoningModel = new MockLanguageModelV3({ doGenerate: async () => ({ - rawCall: { rawPrompt: null, rawSettings: {} }, finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 }, + usage: mockUsage, content: [{ type: "text", text: "Hello, world!" }], warnings: [], }), @@ -34,15 +36,13 @@ export const reasoningModel = new MockLanguageModelV2({ initialDelayInMs: 1000, chunks: getResponseChunksByPrompt(prompt, true), }), - rawCall: { rawPrompt: null, rawSettings: {} }, }), }); -export const titleModel = new MockLanguageModelV2({ +export const titleModel = new MockLanguageModelV3({ doGenerate: async () => ({ - rawCall: { rawPrompt: null, rawSettings: {} }, finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 }, + usage: mockUsage, content: [{ type: "text", text: "This is a test title" }], warnings: [], }), @@ -57,19 +57,17 @@ export const titleModel = new MockLanguageModelV2({ { type: "finish", finishReason: "stop", - usage: { inputTokens: 3, outputTokens: 10, totalTokens: 13 }, + usage: mockUsage, }, ], }), - rawCall: { rawPrompt: null, rawSettings: {} }, }), }); -export const artifactModel = new MockLanguageModelV2({ +export const artifactModel = new MockLanguageModelV3({ doGenerate: async () => ({ - rawCall: { rawPrompt: null, rawSettings: {} }, finishReason: "stop", - usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 }, + usage: mockUsage, content: [{ type: "text", text: "Hello, world!" }], warnings: [], }), @@ -79,6 +77,5 @@ export const artifactModel = new MockLanguageModelV2({ initialDelayInMs: 100, chunks: getResponseChunksByPrompt(prompt), }), - rawCall: { rawPrompt: null, rawSettings: {} }, }), }); diff --git a/lib/ai/tools/get-weather.ts b/lib/ai/tools/get-weather.ts index 3e616e0..ec5d066 100644 --- a/lib/ai/tools/get-weather.ts +++ b/lib/ai/tools/get-weather.ts @@ -40,6 +40,7 @@ export const getWeather = tool({ .describe("City name (e.g., 'San Francisco', 'New York', 'London')") .optional(), }), + needsApproval: true, execute: async (input) => { let latitude: number; let longitude: number; diff --git a/lib/db/migrate.ts b/lib/db/migrate.ts index aec5dcb..e815cb8 100644 --- a/lib/db/migrate.ts +++ b/lib/db/migrate.ts @@ -9,7 +9,8 @@ config({ const runMigrate = async () => { if (!process.env.POSTGRES_URL) { - throw new Error("POSTGRES_URL is not defined"); + console.log("⏭️ POSTGRES_URL not defined, skipping migrations"); + process.exit(0); } const connection = postgres(process.env.POSTGRES_URL, { max: 1 }); diff --git a/lib/db/queries.ts b/lib/db/queries.ts index 47ed0d5..2855423 100644 --- a/lib/db/queries.ts +++ b/lib/db/queries.ts @@ -250,6 +250,20 @@ export async function saveMessages({ messages }: { messages: DBMessage[] }) { } } +export async function updateMessage({ + id, + parts, +}: { + id: string; + parts: DBMessage["parts"]; +}) { + try { + return await db.update(message).set({ parts }).where(eq(message.id, id)); + } catch (_error) { + throw new ChatSDKError("bad_request:database", "Failed to update message"); + } +} + export async function getMessagesByChatId({ id }: { id: string }) { try { return await db diff --git a/lib/utils.ts b/lib/utils.ts index d94b492..7f0c2d5 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -1,6 +1,6 @@ import type { - CoreAssistantMessage, - CoreToolMessage, + AssistantModelMessage, + ToolModelMessage, UIMessage, UIMessagePart, } from 'ai'; @@ -63,7 +63,7 @@ export function generateUUID(): string { }); } -type ResponseMessageWithoutId = CoreToolMessage | CoreAssistantMessage; +type ResponseMessageWithoutId = ToolModelMessage | AssistantModelMessage; type ResponseMessage = ResponseMessageWithoutId & { id: string }; export function getMostRecentUserMessage(messages: UIMessage[]) { diff --git a/package.json b/package.json index 3c77d81..5235cec 100644 --- a/package.json +++ b/package.json @@ -18,9 +18,9 @@ "test": "export PLAYWRIGHT=True && pnpm exec playwright test" }, "dependencies": { - "@ai-sdk/gateway": "^2.0.18", - "@ai-sdk/provider": "2.0.0", - "@ai-sdk/react": "2.0.109", + "@ai-sdk/gateway": "2.0.0-beta.85", + "@ai-sdk/provider": "3.0.0-beta.27", + "@ai-sdk/react": "3.0.0-beta.162", "@codemirror/lang-javascript": "^6.2.2", "@codemirror/lang-python": "^6.1.6", "@codemirror/state": "^6.5.0", @@ -47,7 +47,7 @@ "@vercel/functions": "^2.0.0", "@vercel/otel": "^1.12.0", "@xyflow/react": "^12.10.0", - "ai": "5.0.108", + "ai": "6.0.0-beta.159", "bcrypt-ts": "^5.0.2", "class-variance-authority": "^0.7.1", "classnames": "^2.5.1", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b60080f..07bcac3 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -9,14 +9,14 @@ importers: .: dependencies: '@ai-sdk/gateway': - specifier: ^2.0.18 - version: 2.0.18(zod@3.25.76) + specifier: 2.0.0-beta.85 + version: 2.0.0-beta.85(zod@3.25.76) '@ai-sdk/provider': - specifier: 2.0.0 - version: 2.0.0 + specifier: 3.0.0-beta.27 + version: 3.0.0-beta.27 '@ai-sdk/react': - specifier: 2.0.109 - version: 2.0.109(react@19.0.1)(zod@3.25.76) + specifier: 3.0.0-beta.162 + version: 3.0.0-beta.162(react@19.0.1)(zod@3.25.76) '@codemirror/lang-javascript': specifier: ^6.2.2 version: 6.2.3 @@ -96,8 +96,8 @@ importers: specifier: ^12.10.0 version: 12.10.0(@types/react@18.3.18)(react-dom@19.0.1(react@19.0.1))(react@19.0.1) ai: - specifier: 5.0.108 - version: 5.0.108(zod@3.25.76) + specifier: 6.0.0-beta.159 + version: 6.0.0-beta.159(zod@3.25.76) bcrypt-ts: specifier: ^5.0.2 version: 5.0.3 @@ -303,31 +303,27 @@ importers: packages: - '@ai-sdk/gateway@2.0.18': - resolution: {integrity: sha512-sDQcW+6ck2m0pTIHW6BPHD7S125WD3qNkx/B8sEzJp/hurocmJ5Cni0ybExg6sQMGo+fr/GWOwpHF1cmCdg5rQ==} + '@ai-sdk/gateway@2.0.0-beta.85': + resolution: {integrity: sha512-1LFCTwweCe1KWyBR/v64zbvNJbAu4JPooa+0JVUclcb90RYQH2FDIQCbxN4H6J8is4yaTkHW5tbLNjHpOkxZuA==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 - '@ai-sdk/provider-utils@3.0.18': - resolution: {integrity: sha512-ypv1xXMsgGcNKUP+hglKqtdDuMg68nWHucPPAhIENrbFAI+xCHiqPVN8Zllxyv1TNZwGWUghPxJXU+Mqps0YRQ==} + '@ai-sdk/provider-utils@4.0.0-beta.53': + resolution: {integrity: sha512-83/aNTnKfurb4jdaOSfh1KgxY27SuWZbecc3bUfiUxLMtAMBLusgWMaaSbfl2VHufAoGLXIuRYRq4lXlNCdXzw==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 - '@ai-sdk/provider@2.0.0': - resolution: {integrity: sha512-6o7Y2SeO9vFKB8lArHXehNuusnpddKPk7xqL7T2/b+OvXMRIXUO1rR4wcv1hAFUAT9avGZshty3Wlua/XA7TvA==} + '@ai-sdk/provider@3.0.0-beta.27': + resolution: {integrity: sha512-g/H1lyBQa5TAolD0t9uW552z0dwo2evMPxE9gu22zkEUvQSkOl8F1C0Jg31sUPTn9xmKnDK+hjWsAsZnOFE9mQ==} engines: {node: '>=18'} - '@ai-sdk/react@2.0.109': - resolution: {integrity: sha512-5qM8KuN7bv7E+g6BXkSAYLFjwIfMSTKOA1prjg1zEShJXJyLSc+Yqkd3EfGibm75b7nJAqJNShurDmR/IlQqFQ==} + '@ai-sdk/react@3.0.0-beta.162': + resolution: {integrity: sha512-vgJPUbHJ+y5Ebs1KmCnnjdebP7dsgW4uBEtPvmJJv/5JW7K4nizEFVOoxb2FF/mU7KbbE8qiprnM/XPRgC/znQ==} engines: {node: '>=18'} peerDependencies: - react: ^18 || ^19 || ^19.0.0-rc - zod: ^3.25.76 || ^4.1.8 - peerDependenciesMeta: - zod: - optional: true + react: ^18 || ~19.0.1 || ~19.1.2 || ^19.2.1 '@alloc/quick-lru@5.2.0': resolution: {integrity: sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw==} @@ -2179,8 +2175,8 @@ packages: '@shikijs/vscode-textmate@10.0.2': resolution: {integrity: sha512-83yeghZ2xxin3Nj8z1NMd/NCuca+gsYXswywDy5bHvwlWL8tpTQmzGeUuHd9FC3E/SBEMvzJRwWEOz5gGes9Qg==} - '@standard-schema/spec@1.0.0': - resolution: {integrity: sha512-m2bOd0f2RT9k8QJx1JN85cZYyH1RqFBdlwtkSlf4tBDYLCiiZnv1fIIwacK6cqwXavOydf0NPToMQgpKq+dVlA==} + '@standard-schema/spec@1.1.0': + resolution: {integrity: sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w==} '@swc/helpers@0.5.15': resolution: {integrity: sha512-JQ5TuMi45Owi4/BIMAJBoSQoOJu12oOk/gADqlcUL9JEdHB8vyjUSsxqeNXnmXHjYKMi2WcYtezGEEhqUI/E2g==} @@ -2579,8 +2575,8 @@ packages: engines: {node: '>=0.4.0'} hasBin: true - ai@5.0.108: - resolution: {integrity: sha512-Jex3Lb7V41NNpuqJHKgrwoU6BCLHdI1Pg4qb4GJH4jRIDRXUBySJErHjyN4oTCwbiYCeb/8II9EnqSRPq9EifA==} + ai@6.0.0-beta.159: + resolution: {integrity: sha512-iwyz0iycu0Tu1of9GLGwiXvNgW6dB+hnFE0dzWubTQFKz0C8nB8gjMwRTiNhNXO6HuN9+uokg9DO7HYM2CRExw==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 @@ -4534,33 +4530,33 @@ packages: snapshots: - '@ai-sdk/gateway@2.0.18(zod@3.25.76)': + '@ai-sdk/gateway@2.0.0-beta.85(zod@3.25.76)': dependencies: - '@ai-sdk/provider': 2.0.0 - '@ai-sdk/provider-utils': 3.0.18(zod@3.25.76) + '@ai-sdk/provider': 3.0.0-beta.27 + '@ai-sdk/provider-utils': 4.0.0-beta.53(zod@3.25.76) '@vercel/oidc': 3.0.5 zod: 3.25.76 - '@ai-sdk/provider-utils@3.0.18(zod@3.25.76)': + '@ai-sdk/provider-utils@4.0.0-beta.53(zod@3.25.76)': dependencies: - '@ai-sdk/provider': 2.0.0 - '@standard-schema/spec': 1.0.0 + '@ai-sdk/provider': 3.0.0-beta.27 + '@standard-schema/spec': 1.1.0 eventsource-parser: 3.0.6 zod: 3.25.76 - '@ai-sdk/provider@2.0.0': + '@ai-sdk/provider@3.0.0-beta.27': dependencies: json-schema: 0.4.0 - '@ai-sdk/react@2.0.109(react@19.0.1)(zod@3.25.76)': + '@ai-sdk/react@3.0.0-beta.162(react@19.0.1)(zod@3.25.76)': dependencies: - '@ai-sdk/provider-utils': 3.0.18(zod@3.25.76) - ai: 5.0.108(zod@3.25.76) + '@ai-sdk/provider-utils': 4.0.0-beta.53(zod@3.25.76) + ai: 6.0.0-beta.159(zod@3.25.76) react: 19.0.1 swr: 2.3.3(react@19.0.1) throttleit: 2.1.0 - optionalDependencies: - zod: 3.25.76 + transitivePeerDependencies: + - zod '@alloc/quick-lru@5.2.0': {} @@ -6166,7 +6162,7 @@ snapshots: '@shikijs/vscode-textmate@10.0.2': {} - '@standard-schema/spec@1.0.0': {} + '@standard-schema/spec@1.1.0': {} '@swc/helpers@0.5.15': dependencies: @@ -6576,11 +6572,11 @@ snapshots: acorn@8.15.0: {} - ai@5.0.108(zod@3.25.76): + ai@6.0.0-beta.159(zod@3.25.76): dependencies: - '@ai-sdk/gateway': 2.0.18(zod@3.25.76) - '@ai-sdk/provider': 2.0.0 - '@ai-sdk/provider-utils': 3.0.18(zod@3.25.76) + '@ai-sdk/gateway': 2.0.0-beta.85(zod@3.25.76) + '@ai-sdk/provider': 3.0.0-beta.27 + '@ai-sdk/provider-utils': 4.0.0-beta.53(zod@3.25.76) '@opentelemetry/api': 1.9.0 zod: 3.25.76 diff --git a/tests/prompts/utils.ts b/tests/prompts/utils.ts new file mode 100644 index 0000000..83ccf9d --- /dev/null +++ b/tests/prompts/utils.ts @@ -0,0 +1,30 @@ +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; + +const mockUsage = { + inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 20, text: 20, reasoning: 0 }, +}; + +export function getResponseChunksByPrompt( + _prompt: unknown, + includeReasoning = false +): LanguageModelV3StreamPart[] { + const chunks: LanguageModelV3StreamPart[] = []; + + if (includeReasoning) { + chunks.push( + { type: "reasoning-start", id: "r1" }, + { type: "reasoning-delta", id: "r1", delta: "Let me think about this." }, + { type: "reasoning-end", id: "r1" } + ); + } + + chunks.push( + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "Hello, world!" }, + { type: "text-end", id: "t1" }, + { type: "finish", finishReason: "stop", usage: mockUsage } + ); + + return chunks; +}