fix: preserve message parts order in chat display (#151)

- Fix bug where text after tool calls was merged with initial text
- Group consecutive text/file parts into bubbles while keeping tools in order
- Parts now display as: plan -> tool_result -> additional text
- Remove debug logs from fixToolCallInputs function

Co-authored-by: dayuan.jiang <jiangdy@amazon.co.jp>

app/api/chat/route.ts

@@ -67,41 +67,23 @@ function isMinimalDiagram(xml: string): boolean {
 // Helper function to fix tool call inputs for Bedrock API
 // Bedrock requires toolUse.input to be a JSON object, not a string
 function fixToolCallInputs(messages: any[]): any[] {
-    return messages.map((msg, msgIndex) => {
+    return messages.map((msg) => {
         if (msg.role !== "assistant" || !Array.isArray(msg.content)) {
             return msg
         }
-        const fixedContent = msg.content.map((part: any, partIndex: number) => {
+        const fixedContent = msg.content.map((part: any) => {
             if (part.type === "tool-call") {
-                console.log(
-                    `[fixToolCallInputs] msg[${msgIndex}].content[${partIndex}] tool-call:`,
-                    {
-                        toolName: part.toolName,
-                        inputType: typeof part.input,
-                        input: part.input,
-                    },
-                )
                 if (typeof part.input === "string") {
                     try {
                         const parsed = JSON.parse(part.input)
-                        console.log(
-                            `[fixToolCallInputs] Parsed string input to JSON:`,
-                            parsed,
-                        )
                         return { ...part, input: parsed }
                     } catch {
                         // If parsing fails, wrap the string in an object
-                        console.log(
-                            `[fixToolCallInputs] Failed to parse, wrapping in object`,
-                        )
                         return { ...part, input: { rawInput: part.input } }
                     }
                 }
                 // Input is already an object, but verify it's not null/undefined
                 if (part.input === null || part.input === undefined) {
-                    console.log(
-                        `[fixToolCallInputs] Input is null/undefined, using empty object`,
-                    )
                     return { ...part, input: {} }
                 }
             }

components/chat-message-display.tsx

@@ -509,30 +509,85 @@ export function ChatMessageDisplay({
                                             </div>
                                         </div>
                                     ) : (
-                                        /* Text content in bubble */
+                                        /* Render parts in order, grouping consecutive text/file parts into bubbles */
-                                        message.parts?.some(
+                                        (() => {
-                                            (part) =>
+                                            const parts = message.parts || []
+                                            const groups: {
+                                                type: "content" | "tool"
+                                                parts: typeof parts
+                                                startIndex: number
+                                            }[] = []
+
+                                            parts.forEach((part, index) => {
+                                                const isToolPart =
+                                                    part.type?.startsWith(
+                                                        "tool-",
+                                                    )
+                                                const isContentPart =
                                                     part.type === "text" ||
-                                                part.type === "file",
+                                                    part.type === "file"
-                                        ) && (
+
+                                                if (isToolPart) {
+                                                    groups.push({
+                                                        type: "tool",
+                                                        parts: [part],
+                                                        startIndex: index,
+                                                    })
+                                                } else if (isContentPart) {
+                                                    const lastGroup =
+                                                        groups[
+                                                            groups.length - 1
+                                                        ]
+                                                    if (
+                                                        lastGroup?.type ===
+                                                        "content"
+                                                    ) {
+                                                        lastGroup.parts.push(
+                                                            part,
+                                                        )
+                                                    } else {
+                                                        groups.push({
+                                                            type: "content",
+                                                            parts: [part],
+                                                            startIndex: index,
+                                                        })
+                                                    }
+                                                }
+                                            })
+
+                                            return groups.map(
+                                                (group, groupIndex) => {
+                                                    if (group.type === "tool") {
+                                                        return renderToolPart(
+                                                            group
+                                                                .parts[0] as ToolPartLike,
+                                                        )
+                                                    }
+
+                                                    // Content bubble
+                                                    return (
                                                         <div
+                                                            key={`${message.id}-content-${group.startIndex}`}
                                                             className={`px-4 py-3 text-sm leading-relaxed ${
-                                                    message.role === "user"
+                                                                message.role ===
+                                                                "user"
                                                                     ? "bg-primary text-primary-foreground rounded-2xl rounded-br-md shadow-sm"
                                                                     : message.role ===
                                                                         "system"
                                                                       ? "bg-destructive/10 text-destructive border border-destructive/20 rounded-2xl rounded-bl-md"
                                                                       : "bg-muted/60 text-foreground rounded-2xl rounded-bl-md"
-                                                } ${message.role === "user" && isLastUserMessage && onEditMessage ? "cursor-pointer hover:opacity-90 transition-opacity" : ""}`}
+                                                            } ${message.role === "user" && isLastUserMessage && onEditMessage ? "cursor-pointer hover:opacity-90 transition-opacity" : ""} ${groupIndex > 0 ? "mt-3" : ""}`}
                                                             role={
-                                                    message.role === "user" &&
+                                                                message.role ===
+                                                                    "user" &&
                                                                 isLastUserMessage &&
                                                                 onEditMessage
                                                                     ? "button"
                                                                     : undefined
                                                             }
                                                             tabIndex={
-                                                    message.role === "user" &&
+                                                                message.role ===
+                                                                    "user" &&
                                                                 isLastUserMessage &&
                                                                 onEditMessage
                                                                     ? 0
@@ -555,8 +610,10 @@ export function ChatMessageDisplay({
                                                             }}
                                                             onKeyDown={(e) => {
                                                                 if (
-                                                        (e.key === "Enter" ||
+                                                                    (e.key ===
-                                                            e.key === " ") &&
+                                                                        "Enter" ||
+                                                                        e.key ===
+                                                                            " ") &&
                                                                     message.role ===
                                                                         "user" &&
                                                                     isLastUserMessage &&
@@ -572,20 +629,26 @@ export function ChatMessageDisplay({
                                                                 }
                                                             }}
                                                             title={
-                                                    message.role === "user" &&
+                                                                message.role ===
+                                                                    "user" &&
                                                                 isLastUserMessage &&
                                                                 onEditMessage
                                                                     ? "Click to edit"
                                                                     : undefined
                                                             }
                                                         >
-                                                {message.parts?.map(
+                                                            {group.parts.map(
-                                                    (part, index) => {
+                                                                (
-                                                        switch (part.type) {
+                                                                    part,
-                                                            case "text":
+                                                                    partIndex,
+                                                                ) => {
+                                                                    if (
+                                                                        part.type ===
+                                                                        "text"
+                                                                    ) {
                                                                         return (
                                                                             <div
-                                                                        key={`${message.id}-text-${index}`}
+                                                                                key={`${message.id}-text-${group.startIndex}-${partIndex}`}
                                                                                 className={`prose prose-sm max-w-none break-words [&>*:first-child]:mt-0 [&>*:last-child]:mb-0 ${
                                                                                     message.role ===
                                                                                     "user"
@@ -595,20 +658,34 @@ export function ChatMessageDisplay({
                                                                             >
                                                                                 <ReactMarkdown>
                                                                                     {
-                                                                                part.text
+                                                                                        (
+                                                                                            part as {
+                                                                                                text: string
+                                                                                            }
+                                                                                        )
+                                                                                            .text
                                                                                     }
                                                                                 </ReactMarkdown>
                                                                             </div>
                                                                         )
-                                                            case "file":
+                                                                    }
+                                                                    if (
+                                                                        part.type ===
+                                                                        "file"
+                                                                    ) {
                                                                         return (
                                                                             <div
-                                                                        key={`${message.id}-file-${part.url}`}
+                                                                                key={`${message.id}-file-${group.startIndex}-${partIndex}`}
                                                                                 className="mt-2"
                                                                             >
                                                                                 <Image
                                                                                     src={
-                                                                                part.url
+                                                                                        (
+                                                                                            part as {
+                                                                                                url: string
+                                                                                            }
+                                                                                        )
+                                                                                            .url
                                                                                     }
                                                                                     width={
                                                                                         200
@@ -625,23 +702,16 @@ export function ChatMessageDisplay({
                                                                                 />
                                                                             </div>
                                                                         )
-                                                            default:
-                                                                return null
                                                                     }
+                                                                    return null
                                                                 },
                                                             )}
                                                         </div>
                                                     )
-                                    )}
+                                                },
-                                    {/* Tool calls outside bubble */}
-                                    {message.parts?.map((part) => {
-                                        if (part.type?.startsWith("tool-")) {
-                                            return renderToolPart(
-                                                part as ToolPartLike,
                                             )
-                                        }
+                                        })()
-                                        return null
+                                    )}
-                                    })}
                                     {/* Action buttons for assistant messages */}
                                     {message.role === "assistant" && (
                                         <div className="flex items-center gap-1 mt-2">