fix: finalize artist handoff after chat completion

arpitgupta1214 · arpitgupta1214 · commit cf9f1816b223 · 2026-04-07T00:42:58.000+05:30
diff --git a/lib/chat/__tests__/handleChatCompletion.test.ts b/lib/chat/__tests__/handleChatCompletion.test.ts
@@ -9,6 +9,8 @@ import { sendNewConversationNotification } from "@/lib/telegram/sendNewConversat
 import { generateChatTitle } from "@/lib/chat/generateChatTitle";
 import { handleSendEmailToolOutputs } from "@/lib/emails/handleSendEmailToolOutputs";
 import { sendErrorNotification } from "@/lib/telegram/sendErrorNotification";
+import { copyRoom } from "@/lib/rooms/copyRoom";
+import { copyChatMessages } from "@/lib/chats/copyChatMessages";
 import { handleChatCompletion } from "../handleChatCompletion";
 import type { ChatRequestBody } from "../validateChatRequest";
 
@@ -45,6 +47,14 @@ vi.mock("@/lib/telegram/sendErrorNotification", () => ({
   sendErrorNotification: vi.fn(),
 }));
 
+vi.mock("@/lib/rooms/copyRoom", () => ({
+  copyRoom: vi.fn(),
+}));
+
+vi.mock("@/lib/chats/copyChatMessages", () => ({
+  copyChatMessages: vi.fn(),
+}));
+
 const mockSelectAccountEmails = vi.mocked(selectAccountEmails);
 const mockSelectRoom = vi.mocked(selectRoom);
 const mockUpsertRoom = vi.mocked(upsertRoom);
@@ -53,6 +63,8 @@ const mockSendNewConversationNotification = vi.mocked(sendNewConversationNotific
 const mockGenerateChatTitle = vi.mocked(generateChatTitle);
 const mockHandleSendEmailToolOutputs = vi.mocked(handleSendEmailToolOutputs);
 const mockSendErrorNotification = vi.mocked(sendErrorNotification);
+const mockCopyRoom = vi.mocked(copyRoom);
+const mockCopyChatMessages = vi.mocked(copyChatMessages);
 
 // Helper to create mock UIMessage
 /**
@@ -92,6 +104,12 @@ describe("handleChatCompletion", () => {
     mockSelectRoom.mockResolvedValue({ id: "room-456" });
     mockUpsertMemory.mockResolvedValue(null);
     mockHandleSendEmailToolOutputs.mockResolvedValue();
+    mockCopyRoom.mockResolvedValue("new-room-123");
+    mockCopyChatMessages.mockResolvedValue({
+      success: true,
+      copiedCount: 2,
+      clearedExisting: true,
+    });
   });
 
   afterEach(() => {
@@ -267,7 +285,7 @@ describe("handleChatCompletion", () => {
       const responseMessages = [createMockUIMessage("resp-1", "assistant", "Hi!")];
 
       // Should not throw
-      await expect(handleChatCompletion(body, responseMessages)).resolves.toBeUndefined();
+      await expect(handleChatCompletion(body, responseMessages)).resolves.toEqual({});
     });
   });
 
@@ -317,5 +335,42 @@ describe("handleChatCompletion", () => {
         }),
       );
     });
+
+    it("returns redirect path after creating and copying the final artist room", async () => {
+      const body = createMockBody();
+      const responseMessages: UIMessage[] = [
+        {
+          id: "resp-1",
+          role: "assistant",
+          parts: [
+            {
+              type: "tool-create_new_artist",
+              toolCallId: "tool-1",
+              state: "output-available",
+              input: {},
+              output: {
+                artist: {
+                  account_id: "artist-123",
+                  name: "Test Artist",
+                },
+                artistAccountId: "artist-123",
+                message: "ok",
+              },
+            } as any,
+          ],
+          createdAt: new Date(),
+        },
+      ];
+
+      const result = await handleChatCompletion(body, responseMessages);
+
+      expect(mockCopyRoom).toHaveBeenCalledWith("room-456", "artist-123");
+      expect(mockCopyChatMessages).toHaveBeenCalledWith({
+        sourceChatId: "room-456",
+        targetChatId: "new-room-123",
+        clearExisting: true,
+      });
+      expect(result).toEqual({ redirectPath: "/chat/new-room-123" });
+    });
   });
 });
diff --git a/lib/chat/__tests__/handleChatStream.test.ts b/lib/chat/__tests__/handleChatStream.test.ts
@@ -5,6 +5,7 @@ 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 { handleChatCompletion } from "@/lib/chat/handleChatCompletion";
 import { createUIMessageStream, createUIMessageStreamResponse } from "ai";
 import { handleChatStream } from "../handleChatStream";
 
@@ -71,6 +72,7 @@ const mockGetApiKeyAccountId = vi.mocked(getApiKeyAccountId);
 const mockValidateOverrideAccountId = vi.mocked(validateOverrideAccountId);
 const mockSetupConversation = vi.mocked(setupConversation);
 const mockSetupChatRequest = vi.mocked(setupChatRequest);
+const mockHandleChatCompletion = vi.mocked(handleChatCompletion);
 const mockCreateUIMessageStream = vi.mocked(createUIMessageStream);
 const mockCreateUIMessageStreamResponse = vi.mocked(createUIMessageStreamResponse);
 
@@ -99,6 +101,7 @@ describe("handleChatStream", () => {
       roomId: roomId || "mock-room-id",
       memoryId: "mock-memory-id",
     }));
+    mockHandleChatCompletion.mockResolvedValue({});
   });
 
   afterEach(() => {
@@ -247,6 +250,68 @@ describe("handleChatStream", () => {
         }),
       );
     });
+
+    it("uses sendFinish false and emits redirect data after completion", async () => {
+      mockGetApiKeyAccountId.mockResolvedValue("account-123");
+      mockHandleChatCompletion.mockResolvedValue({ redirectPath: "/chat/new-room-123" });
+
+      const toUIMessageStream = vi.fn().mockReturnValue(new ReadableStream());
+      const mockAgent = {
+        stream: vi.fn().mockResolvedValue({
+          toUIMessageStream,
+          usage: Promise.resolve({ inputTokens: 100, outputTokens: 50 }),
+        }),
+        tools: {},
+      };
+
+      mockSetupChatRequest.mockResolvedValue({
+        agent: mockAgent,
+        messages: [],
+      } as any);
+
+      const mockStream = new ReadableStream();
+      mockCreateUIMessageStream.mockReturnValue(mockStream);
+      mockCreateUIMessageStreamResponse.mockReturnValue(new Response(mockStream));
+
+      const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "valid-key" });
+
+      await handleChatStream(request as any);
+
+      const execute = mockCreateUIMessageStream.mock.calls[0][0].execute;
+      const writer = {
+        merge: vi.fn(),
+        write: vi.fn(),
+        onError: undefined,
+      };
+
+      await execute({ writer } as any);
+
+      expect(toUIMessageStream).toHaveBeenCalledWith(
+        expect.objectContaining({
+          sendFinish: false,
+          onFinish: expect.any(Function),
+        }),
+      );
+
+      const onFinish = toUIMessageStream.mock.calls[0][0].onFinish;
+      await onFinish({
+        isAborted: false,
+        finishReason: "stop",
+        messages: [{ id: "resp-1", role: "assistant", parts: [] }],
+        responseMessage: { id: "resp-1", role: "assistant", parts: [] },
+        isContinuation: false,
+      });
+
+      expect(writer.write).toHaveBeenCalledWith({
+        type: "data-redirect",
+        data: { path: "/chat/new-room-123" },
+        transient: true,
+      });
+      expect(writer.write).toHaveBeenCalledWith({
+        type: "finish",
+        finishReason: "stop",
+      });
+    });
   });
 
   describe("error handling", () => {
diff --git a/lib/chat/handleChatCompletion.ts b/lib/chat/handleChatCompletion.ts
@@ -1,4 +1,8 @@
-import type { UIMessage } from "ai";
+import {
+  getToolOrDynamicToolName,
+  isToolOrDynamicToolUIPart,
+  type UIMessage,
+} from "ai";
 import selectAccountEmails from "@/lib/supabase/account_emails/selectAccountEmails";
 import selectRoom from "@/lib/supabase/rooms/selectRoom";
 import { upsertRoom } from "@/lib/supabase/rooms/upsertRoom";
@@ -11,6 +15,38 @@ import { handleSendEmailToolOutputs } from "@/lib/emails/handleSendEmailToolOutp
 import { sendErrorNotification } from "@/lib/telegram/sendErrorNotification";
 import { serializeError } from "@/lib/errors/serializeError";
 import type { ChatRequestBody } from "./validateChatRequest";
+import { copyRoom } from "@/lib/rooms/copyRoom";
+import { copyChatMessages } from "@/lib/chats/copyChatMessages";
+import type { CreateNewArtistResult } from "@/lib/mcp/tools/artists/registerCreateNewArtistTool";
+
+export interface ChatCompletionResult {
+  redirectPath?: string;
+}
+
+function getCreateArtistResult(responseMessages: UIMessage[]): CreateNewArtistResult | null {
+  for (const message of responseMessages) {
+    for (const part of message.parts) {
+      if (!isToolOrDynamicToolUIPart(part)) continue;
+      if (getToolOrDynamicToolName(part) !== "create_new_artist") continue;
+      if (part.state !== "output-available") continue;
+
+      if (part.type === "dynamic-tool") {
+        const text = part.output?.content?.[0]?.text;
+        if (!text) continue;
+
+        try {
+          return JSON.parse(text) as CreateNewArtistResult;
+        } catch {
+          continue;
+        }
+      }
+
+      return part.output as CreateNewArtistResult;
+    }
+  }
+
+  return null;
+}
 
 /**
  * Handles post-chat-completion tasks:
@@ -28,7 +64,7 @@ import type { ChatRequestBody } from "./validateChatRequest";
 export async function handleChatCompletion(
   body: ChatRequestBody,
   responseMessages: UIMessage[],
-): Promise<void> {
+): Promise<ChatCompletionResult> {
   try {
     const { messages, roomId = "", accountId, artistId } = body;
 
@@ -81,14 +117,39 @@ export async function handleChatCompletion(
       content: filterMessageContentForMemories(responseMessages[responseMessages.length - 1]),
     });
 
+    let redirectPath: string | undefined;
+    const createArtistResult = getCreateArtistResult(responseMessages);
+    if (createArtistResult?.artistAccountId) {
+      const newRoomId = await copyRoom(roomId, createArtistResult.artistAccountId);
+
+      if (newRoomId) {
+        const copyResult = await copyChatMessages({
+          sourceChatId: roomId,
+          targetChatId: newRoomId,
+          clearExisting: true,
+        });
+
+        if (copyResult.success) {
+          redirectPath = `/chat/${newRoomId}`;
+        } else {
+          console.error("Failed to copy final artist conversation:", copyResult.error);
+        }
+      } else {
+        console.error("Failed to create final artist conversation room");
+      }
+    }
+
     // Process any email tool outputs
     await handleSendEmailToolOutputs(responseMessages);
+
+    return { redirectPath };
   } catch (error) {
     sendErrorNotification({
       ...body,
       path: "/api/chat",
       error: serializeError(error),
     });
     console.error("Failed to save chat", error);
+    return {};
   }
 }
diff --git a/lib/chat/handleChatStream.ts b/lib/chat/handleChatStream.ts
@@ -38,23 +38,42 @@ export async function handleChatStream(request: NextRequest): Promise<Response>
       execute: async options => {
         const { writer } = options;
         streamResult = await agent.stream(chatConfig);
-        writer.merge(streamResult.toUIMessageStream());
-      },
-      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);
-        if (streamResult) {
-          await handleChatCredits({
-            usage: await streamResult.usage,
-            model: body.model ?? DEFAULT_MODEL,
-            accountId: body.accountId,
-          });
-        }
+        writer.merge(
+          streamResult.toUIMessageStream({
+            sendFinish: false,
+            onFinish: async event => {
+              if (event.isAborted) {
+                return;
+              }
+
+              const assistantMessages = event.messages.filter(
+                message => message.role === "assistant",
+              );
+              const responseMessages =
+                assistantMessages.length > 0 ? assistantMessages : [event.responseMessage];
+              const { redirectPath } = await handleChatCompletion(body, responseMessages);
+
+              if (redirectPath) {
+                writer.write({
+                  type: "data-redirect",
+                  data: { path: redirectPath },
+                  transient: true,
+                });
+              }
+
+              writer.write({
+                type: "finish",
+                finishReason: event.finishReason,
+              });
+
+              await handleChatCredits({
+                usage: await streamResult!.usage,
+                model: body.model ?? DEFAULT_MODEL,
+                accountId: body.accountId,
+              });
+            },
+          }),
+        );
       },
       onError: e => {
         console.error("/api/chat onError:", e);
diff --git a/lib/mcp/tools/artists/__tests__/registerCreateNewArtistTool.test.ts b/lib/mcp/tools/artists/__tests__/registerCreateNewArtistTool.test.ts
diff --git a/lib/mcp/tools/artists/registerCreateNewArtistTool.ts b/lib/mcp/tools/artists/registerCreateNewArtistTool.ts
