From 96aff7992b2c1864b8a9504deb0c11e855b157f7 Mon Sep 17 00:00:00 2001 From: Sweets Sweetman Date: Mon, 19 Jan 2026 13:39:45 -0500 Subject: [PATCH 1/2] feat: use saveChatCompletion in handleChatStream to match handleChatGenerate (MYC-3962) Replace handleChatCompletion with saveChatCompletion in handleChatStream to: - Match the message persistence behavior of handleChatGenerate - Remove duplicate room creation (now handled in validateChatRequest) - Remove duplicate user message persistence (now handled in setupConversation) - Keep assistant message persistence with proper error handling Add tests for: - saveChatCompletion called with text from last assistant message - Fallback to responseMessage when no assistant messages - No call when stream is aborted - Error logging without throwing Co-Authored-By: Claude Opus 4.5 --- lib/chat/__tests__/handleChatStream.test.ts | 281 +++++++++++++++++--- lib/chat/handleChatStream.ts | 33 ++- 2 files changed, 272 insertions(+), 42 deletions(-) diff --git a/lib/chat/__tests__/handleChatStream.test.ts b/lib/chat/__tests__/handleChatStream.test.ts index 5943feed..41fc7ceb 100644 --- a/lib/chat/__tests__/handleChatStream.test.ts +++ b/lib/chat/__tests__/handleChatStream.test.ts @@ -1,6 +1,14 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { NextResponse } from "next/server"; +import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; +import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; +import { setupChatRequest } from "@/lib/chat/setupChatRequest"; +import { setupConversation } from "@/lib/chat/setupConversation"; +import { saveChatCompletion } from "@/lib/chat/saveChatCompletion"; +import { createUIMessageStream, createUIMessageStreamResponse } from "ai"; +import { handleChatStream } from "../handleChatStream"; + // Mock all dependencies before importing the module under test vi.mock("@/lib/auth/getApiKeyAccountId", () => ({ getApiKeyAccountId: vi.fn(), @@ -23,26 +31,28 @@ vi.mock("@/lib/organizations/validateOrganizationAccess", () => ({ })); vi.mock("@/lib/chat/setupConversation", () => ({ - setupConversation: vi.fn().mockResolvedValue({ roomId: "mock-room-id", memoryId: "mock-memory-id" }), + setupConversation: vi + .fn() + .mockResolvedValue({ roomId: "mock-room-id", memoryId: "mock-memory-id" }), })); vi.mock("@/lib/chat/validateMessages", () => ({ - validateMessages: vi.fn((messages) => ({ + validateMessages: vi.fn(messages => ({ lastMessage: messages[messages.length - 1] || { id: "mock-id", role: "user", parts: [] }, validMessages: messages, })), })); vi.mock("@/lib/messages/convertToUiMessages", () => ({ - default: vi.fn((messages) => messages), + default: vi.fn(messages => messages), })); vi.mock("@/lib/chat/setupChatRequest", () => ({ setupChatRequest: vi.fn(), })); -vi.mock("@/lib/chat/handleChatCompletion", () => ({ - handleChatCompletion: vi.fn(), +vi.mock("@/lib/chat/saveChatCompletion", () => ({ + saveChatCompletion: vi.fn(), })); vi.mock("ai", () => ({ @@ -50,25 +60,21 @@ vi.mock("ai", () => ({ createUIMessageStreamResponse: vi.fn(), })); -import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; -import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; -import { setupChatRequest } from "@/lib/chat/setupChatRequest"; -import { setupConversation } from "@/lib/chat/setupConversation"; -import { createUIMessageStream, createUIMessageStreamResponse } from "ai"; -import { handleChatStream } from "../handleChatStream"; - const mockGetApiKeyAccountId = vi.mocked(getApiKeyAccountId); const mockValidateOverrideAccountId = vi.mocked(validateOverrideAccountId); const mockSetupConversation = vi.mocked(setupConversation); const mockSetupChatRequest = vi.mocked(setupChatRequest); +const mockSaveChatCompletion = vi.mocked(saveChatCompletion); const mockCreateUIMessageStream = vi.mocked(createUIMessageStream); const mockCreateUIMessageStreamResponse = vi.mocked(createUIMessageStreamResponse); // Helper to create mock NextRequest -function createMockRequest( - body: unknown, - headers: Record = {}, -): Request { +/** + * + * @param body + * @param headers + */ +function createMockRequest(body: unknown, headers: Record = {}): Request { return { json: () => Promise.resolve(body), headers: { @@ -97,10 +103,7 @@ describe("handleChatStream", () => { it("returns 400 error when neither messages nor prompt is provided", async () => { mockGetApiKeyAccountId.mockResolvedValue("account-123"); - const request = createMockRequest( - { roomId: "room-123" }, - { "x-api-key": "test-key" }, - ); + const request = createMockRequest({ roomId: "room-123" }, { "x-api-key": "test-key" }); const result = await handleChatStream(request as any); @@ -151,10 +154,7 @@ describe("handleChatStream", () => { const mockResponse = new Response(mockStream); mockCreateUIMessageStreamResponse.mockReturnValue(mockResponse); - const request = createMockRequest( - { prompt: "Hello, world!" }, - { "x-api-key": "valid-key" }, - ); + const request = createMockRequest({ prompt: "Hello, world!" }, { "x-api-key": "valid-key" }); const result = await handleChatStream(request as any); @@ -193,10 +193,7 @@ describe("handleChatStream", () => { mockCreateUIMessageStreamResponse.mockReturnValue(new Response(mockStream)); const messages = [{ role: "user", content: "Hello" }]; - const request = createMockRequest( - { messages }, - { "x-api-key": "valid-key" }, - ); + const request = createMockRequest({ messages }, { "x-api-key": "valid-key" }); await handleChatStream(request as any); @@ -263,10 +260,7 @@ describe("handleChatStream", () => { mockGetApiKeyAccountId.mockResolvedValue("account-123"); mockSetupChatRequest.mockRejectedValue(new Error("Setup failed")); - const request = createMockRequest( - { prompt: "Hello" }, - { "x-api-key": "valid-key" }, - ); + const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "valid-key" }); const result = await handleChatStream(request as any); @@ -321,4 +315,227 @@ describe("handleChatStream", () => { ); }); }); + + describe("message persistence", () => { + it("calls saveChatCompletion with text from last assistant message in onFinish", async () => { + mockGetApiKeyAccountId.mockResolvedValue("account-123"); + + const mockAgent = { + stream: vi.fn().mockResolvedValue({ + toUIMessageStream: vi.fn().mockReturnValue(new ReadableStream()), + usage: Promise.resolve({ inputTokens: 100, outputTokens: 50 }), + }), + tools: {}, + }; + + mockSetupChatRequest.mockResolvedValue({ + agent: mockAgent, + model: "gpt-4", + instructions: "You are a helpful assistant", + system: "You are a helpful assistant", + messages: [], + experimental_generateMessageId: vi.fn(), + tools: {}, + providerOptions: {}, + } as any); + + // Capture the onFinish callback + let capturedOnFinish: ((event: any) => Promise) | undefined; + mockCreateUIMessageStream.mockImplementation((options: any) => { + capturedOnFinish = options.onFinish; + return new ReadableStream(); + }); + + mockCreateUIMessageStreamResponse.mockReturnValue(new Response(new ReadableStream())); + + const request = createMockRequest( + { prompt: "Hello", roomId: "test-room-id" }, + { "x-api-key": "valid-key" }, + ); + + await handleChatStream(request as any); + + // Simulate onFinish being called with assistant messages + expect(capturedOnFinish).toBeDefined(); + await capturedOnFinish!({ + isAborted: false, + messages: [ + { + id: "msg-1", + role: "assistant", + parts: [{ type: "text", text: "Hello! How can I help you?" }], + }, + ], + responseMessage: { + id: "msg-fallback", + role: "assistant", + parts: [{ type: "text", text: "Fallback response" }], + }, + }); + + expect(mockSaveChatCompletion).toHaveBeenCalledWith({ + text: "Hello! How can I help you?", + roomId: "test-room-id", + }); + }); + + it("uses responseMessage as fallback when no assistant messages", async () => { + mockGetApiKeyAccountId.mockResolvedValue("account-123"); + + const mockAgent = { + stream: vi.fn().mockResolvedValue({ + toUIMessageStream: vi.fn().mockReturnValue(new ReadableStream()), + usage: Promise.resolve({ inputTokens: 100, outputTokens: 50 }), + }), + tools: {}, + }; + + mockSetupChatRequest.mockResolvedValue({ + agent: mockAgent, + model: "gpt-4", + instructions: "test", + system: "test", + messages: [], + experimental_generateMessageId: vi.fn(), + tools: {}, + providerOptions: {}, + } as any); + + let capturedOnFinish: ((event: any) => Promise) | undefined; + mockCreateUIMessageStream.mockImplementation((options: any) => { + capturedOnFinish = options.onFinish; + return new ReadableStream(); + }); + + mockCreateUIMessageStreamResponse.mockReturnValue(new Response(new ReadableStream())); + + const request = createMockRequest( + { prompt: "Hello", roomId: "test-room-id" }, + { "x-api-key": "valid-key" }, + ); + + await handleChatStream(request as any); + + await capturedOnFinish!({ + isAborted: false, + messages: [], // No assistant messages + responseMessage: { + id: "msg-fallback", + role: "assistant", + parts: [{ type: "text", text: "Fallback response" }], + }, + }); + + expect(mockSaveChatCompletion).toHaveBeenCalledWith({ + text: "Fallback response", + roomId: "test-room-id", + }); + }); + + it("does not call saveChatCompletion when stream is aborted", async () => { + mockGetApiKeyAccountId.mockResolvedValue("account-123"); + + const mockAgent = { + stream: vi.fn().mockResolvedValue({ + toUIMessageStream: vi.fn().mockReturnValue(new ReadableStream()), + usage: Promise.resolve({ inputTokens: 100, outputTokens: 50 }), + }), + tools: {}, + }; + + mockSetupChatRequest.mockResolvedValue({ + agent: mockAgent, + model: "gpt-4", + instructions: "test", + system: "test", + messages: [], + experimental_generateMessageId: vi.fn(), + tools: {}, + providerOptions: {}, + } as any); + + let capturedOnFinish: ((event: any) => Promise) | undefined; + mockCreateUIMessageStream.mockImplementation((options: any) => { + capturedOnFinish = options.onFinish; + return new ReadableStream(); + }); + + mockCreateUIMessageStreamResponse.mockReturnValue(new Response(new ReadableStream())); + + const request = createMockRequest( + { prompt: "Hello", roomId: "test-room-id" }, + { "x-api-key": "valid-key" }, + ); + + await handleChatStream(request as any); + + await capturedOnFinish!({ + isAborted: true, + messages: [], + responseMessage: null, + }); + + expect(mockSaveChatCompletion).not.toHaveBeenCalled(); + }); + + it("logs error but does not throw when saveChatCompletion fails", async () => { + mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockSaveChatCompletion.mockRejectedValue(new Error("Database error")); + + const mockAgent = { + stream: vi.fn().mockResolvedValue({ + toUIMessageStream: vi.fn().mockReturnValue(new ReadableStream()), + usage: Promise.resolve({ inputTokens: 100, outputTokens: 50 }), + }), + tools: {}, + }; + + mockSetupChatRequest.mockResolvedValue({ + agent: mockAgent, + model: "gpt-4", + instructions: "test", + system: "test", + messages: [], + experimental_generateMessageId: vi.fn(), + tools: {}, + providerOptions: {}, + } as any); + + let capturedOnFinish: ((event: any) => Promise) | undefined; + mockCreateUIMessageStream.mockImplementation((options: any) => { + capturedOnFinish = options.onFinish; + return new ReadableStream(); + }); + + mockCreateUIMessageStreamResponse.mockReturnValue(new Response(new ReadableStream())); + + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + const request = createMockRequest( + { prompt: "Hello", roomId: "test-room-id" }, + { "x-api-key": "valid-key" }, + ); + + await handleChatStream(request as any); + + // Should not throw + await capturedOnFinish!({ + isAborted: false, + messages: [ + { + id: "msg-1", + role: "assistant", + parts: [{ type: "text", text: "Hello!" }], + }, + ], + responseMessage: null, + }); + + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Failed to persist assistant message:", + expect.any(Error), + ); + consoleErrorSpy.mockRestore(); + }); + }); }); diff --git a/lib/chat/handleChatStream.ts b/lib/chat/handleChatStream.ts index 56cfb963..32277433 100644 --- a/lib/chat/handleChatStream.ts +++ b/lib/chat/handleChatStream.ts @@ -1,6 +1,6 @@ import { NextRequest, NextResponse } from "next/server"; import { createUIMessageStream, createUIMessageStreamResponse } from "ai"; -import { handleChatCompletion } from "./handleChatCompletion"; +import { saveChatCompletion } from "./saveChatCompletion"; import { validateChatRequest } from "./validateChatRequest"; import { setupChatRequest } from "./setupChatRequest"; import { getCorsHeaders } from "@/lib/networking/getCorsHeaders"; @@ -31,25 +31,38 @@ export async function handleChatStream(request: NextRequest): Promise const stream = createUIMessageStream({ originalMessages: body.messages, generateId: generateUUID, - execute: async (options) => { + execute: async options => { const { writer } = options; const result = await agent.stream(chatConfig); writer.merge(result.toUIMessageStream()); // Note: Credit handling and chat completion handling will be added // as part of the handleChatCredits and handleChatCompletion migrations }, - onFinish: async (event) => { + onFinish: async event => { if (event.isAborted) { return; } - const assistantMessages = event.messages.filter( - (message) => message.role === "assistant", - ); - const responseMessages = - assistantMessages.length > 0 ? assistantMessages : [event.responseMessage]; - await handleChatCompletion(body, responseMessages); + const assistantMessages = event.messages.filter(message => message.role === "assistant"); + const lastAssistantMessage = + assistantMessages.length > 0 + ? assistantMessages[assistantMessages.length - 1] + : event.responseMessage; + + // Extract text from the assistant message + const text = lastAssistantMessage.parts.find(part => part.type === "text")?.text || ""; + + // Save assistant message to database (matches handleChatGenerate behavior) + try { + await saveChatCompletion({ + text, + roomId: body.roomId, + }); + } catch (error) { + // Log error but don't fail the request - message persistence is non-critical + console.error("Failed to persist assistant message:", error); + } }, - onError: (e) => { + onError: e => { console.error("/api/chat onError:", e); return JSON.stringify({ status: "error", From 6460c843db7c8f72c95ff67a02ed2076713d51c5 Mon Sep 17 00:00:00 2001 From: Sweets Sweetman Date: Mon, 19 Jan 2026 13:46:18 -0500 Subject: [PATCH 2/2] refactor: extract getUiMessageText utility for DRY text extraction (MYC-3962) - Create getUiMessageText utility to extract text from UIMessage parts - Update handleChatStream to use getUiMessageText instead of inline logic - Update getLatestUserMessageText to use getUiMessageText for consistency - Add 6 unit tests for getUiMessageText Co-Authored-By: Claude Opus 4.5 --- lib/chat/handleChatStream.ts | 4 +- .../__tests__/getUiMessageText.test.ts | 68 +++++++++++++++++++ lib/messages/getLatestUserMessageText.ts | 6 +- lib/messages/getUiMessageText.ts | 11 +++ 4 files changed, 85 insertions(+), 4 deletions(-) create mode 100644 lib/messages/__tests__/getUiMessageText.test.ts create mode 100644 lib/messages/getUiMessageText.ts diff --git a/lib/chat/handleChatStream.ts b/lib/chat/handleChatStream.ts index 32277433..ff962bd5 100644 --- a/lib/chat/handleChatStream.ts +++ b/lib/chat/handleChatStream.ts @@ -5,6 +5,7 @@ import { validateChatRequest } from "./validateChatRequest"; import { setupChatRequest } from "./setupChatRequest"; import { getCorsHeaders } from "@/lib/networking/getCorsHeaders"; import generateUUID from "@/lib/uuid/generateUUID"; +import getUiMessageText from "@/lib/messages/getUiMessageText"; /** * Handles a streaming chat request. @@ -48,8 +49,7 @@ export async function handleChatStream(request: NextRequest): Promise ? assistantMessages[assistantMessages.length - 1] : event.responseMessage; - // Extract text from the assistant message - const text = lastAssistantMessage.parts.find(part => part.type === "text")?.text || ""; + const text = getUiMessageText(lastAssistantMessage); // Save assistant message to database (matches handleChatGenerate behavior) try { diff --git a/lib/messages/__tests__/getUiMessageText.test.ts b/lib/messages/__tests__/getUiMessageText.test.ts new file mode 100644 index 00000000..f6ce58ba --- /dev/null +++ b/lib/messages/__tests__/getUiMessageText.test.ts @@ -0,0 +1,68 @@ +import { describe, it, expect } from "vitest"; +import { UIMessage } from "ai"; +import getUiMessageText from "../getUiMessageText"; + +describe("getUiMessageText", () => { + it("extracts text from a UIMessage with text part", () => { + const message: UIMessage = { + id: "msg-1", + role: "assistant", + parts: [{ type: "text", text: "Hello, world!" }], + }; + + expect(getUiMessageText(message)).toBe("Hello, world!"); + }); + + it("returns first text part when multiple parts exist", () => { + const message: UIMessage = { + id: "msg-2", + role: "assistant", + parts: [ + { type: "text", text: "First text" }, + { type: "text", text: "Second text" }, + ], + }; + + expect(getUiMessageText(message)).toBe("First text"); + }); + + it("returns empty string when no text parts exist", () => { + const message: UIMessage = { + id: "msg-3", + role: "assistant", + parts: [], + }; + + expect(getUiMessageText(message)).toBe(""); + }); + + it("returns empty string when parts array only has non-text parts", () => { + const message = { + id: "msg-4", + role: "assistant", + parts: [{ type: "tool-invocation", toolInvocationId: "123", toolName: "test", state: "result" }], + } as UIMessage; + + expect(getUiMessageText(message)).toBe(""); + }); + + it("works with user messages", () => { + const message: UIMessage = { + id: "msg-5", + role: "user", + parts: [{ type: "text", text: "User question" }], + }; + + expect(getUiMessageText(message)).toBe("User question"); + }); + + it("works with system messages", () => { + const message: UIMessage = { + id: "msg-6", + role: "system", + parts: [{ type: "text", text: "System prompt" }], + }; + + expect(getUiMessageText(message)).toBe("System prompt"); + }); +}); diff --git a/lib/messages/getLatestUserMessageText.ts b/lib/messages/getLatestUserMessageText.ts index 2d83f506..c52f246c 100644 --- a/lib/messages/getLatestUserMessageText.ts +++ b/lib/messages/getLatestUserMessageText.ts @@ -1,4 +1,5 @@ import { UIMessage } from "ai"; +import getUiMessageText from "./getUiMessageText"; /** * Extracts the text content from the most recent user message @@ -7,7 +8,8 @@ import { UIMessage } from "ai"; * @returns The text content of the latest user message, or empty string if none found */ export default function getLatestUserMessageText(messages: UIMessage[]): string { - const userMessages = messages.filter((msg) => msg.role === "user"); + const userMessages = messages.filter(msg => msg.role === "user"); const latestUserMessage = userMessages[userMessages.length - 1]; - return latestUserMessage?.parts?.find((part) => part.type === "text")?.text || ""; + if (!latestUserMessage) return ""; + return getUiMessageText(latestUserMessage); } diff --git a/lib/messages/getUiMessageText.ts b/lib/messages/getUiMessageText.ts new file mode 100644 index 00000000..c8a4a34b --- /dev/null +++ b/lib/messages/getUiMessageText.ts @@ -0,0 +1,11 @@ +import { UIMessage } from "ai"; + +/** + * Extracts text content from a UIMessage. + * + * @param message - The UIMessage to extract text from + * @returns The text content of the first text part, or empty string if none found + */ +export default function getUiMessageText(message: UIMessage): string { + return message.parts?.find((part) => part.type === "text")?.text || ""; +}