diff --git a/apps/dbagent/src/components/chats/chats-ui.tsx b/apps/dbagent/src/components/chats/chats-ui.tsx
index 24562830..0bafe4fd 100644
--- a/apps/dbagent/src/components/chats/chats-ui.tsx
+++ b/apps/dbagent/src/components/chats/chats-ui.tsx
@@ -18,7 +18,7 @@ import {
   Textarea
 } from '@internal/components';
 import { useChat } from 'ai/react';
-import { Bot, Clock, Send, User } from 'lucide-react';
+import { Bot, Clock, Lightbulb, Send, User, Wrench } from 'lucide-react';
 import { useSearchParams } from 'next/navigation';
 import { useEffect, useRef, useState } from 'react';
 import ReactMarkdown from 'react-markdown';
@@ -159,44 +159,66 @@ export function ChatsUI({ connections }: { connections: DbConnection[] }) {
                 </div>
               )}
               {messages.map((message) => (
-                <div
-                  key={message.id}
-                  className={`flex gap-3 ${message.role === 'assistant' ? 'flex-row' : 'flex-row-reverse'}`}
-                >
-                  <Avatar>
-                    <AvatarFallback>
-                      {message.role === 'assistant' ? <Bot className="h-6 w-6" /> : <User className="h-6 w-6" />}
-                    </AvatarFallback>
-                  </Avatar>
-                  <div
-                    className={`max-w-[80%] rounded-lg px-4 py-2 ${
-                      message.role === 'assistant' ? 'bg-secondary' : 'bg-primary text-primary-foreground ml-auto'
-                    }`}
-                  >
-                    <ReactMarkdown
-                      components={{
-                        p: ({ children }) => <p className="mb-2 text-sm last:mb-0">{children}</p>,
-                        ul: ({ children }) => <ul className="mb-2 list-inside list-disc text-sm">{children}</ul>,
-                        ol: ({ children }) => <ol className="mb-2 list-inside list-decimal text-sm">{children}</ol>,
-                        li: ({ children }) => <li className="mb-1">{children}</li>,
-                        h1: ({ children }) => <h1 className="mb-2 text-xl font-bold">{children}</h1>,
-                        h2: ({ children }) => <h2 className="mb-2 text-lg font-semibold">{children}</h2>,
-                        h3: ({ children }) => <h3 className="mb-2 text-base font-medium">{children}</h3>,
-                        code: ({ children }) => (
-                          <code className="rounded bg-black px-1 py-0.5 font-mono text-sm text-white">{children}</code>
-                        )
-                      }}
+                <div key={message.id} className="space-y-4">
+                  {message.parts.map((part, index) => (
+                    <div
+                      key={`${message.id}-${index}`}
+                      className={`flex gap-3 ${message.role === 'assistant' ? 'flex-row' : 'flex-row-reverse'}`}
                     >
-                      {message.content}
-                    </ReactMarkdown>
-                    {message.toolInvocations &&
-                      message.toolInvocations.map((tool, index) => (
-                        <div key={index} className="text-muted-foreground mt-2 text-xs">
-                          <Clock className="mr-1 inline-block h-4 w-4" />
-                          Tool called: {tool.toolName}
-                        </div>
-                      ))}
-                  </div>
+                      <Avatar>
+                        <AvatarFallback>
+                          {part.type === 'text' && message.role === 'user' ? (
+                            <User className="h-6 w-6" />
+                          ) : part.type === 'text' ? (
+                            <Bot className="h-6 w-6" />
+                          ) : part.type === 'tool-invocation' ? (
+                            <Wrench className="h-6 w-6" />
+                          ) : part.type === 'reasoning' ? (
+                            <Lightbulb className="h-6 w-6" />
+                          ) : null}
+                        </AvatarFallback>
+                      </Avatar>
+
+                      <div
+                        className={`max-w-[80%] rounded-lg px-4 py-2 ${
+                          message.role === 'assistant' ? 'bg-secondary' : 'bg-primary text-primary-foreground ml-auto'
+                        }`}
+                      >
+                        {part.type === 'text' ? (
+                          <ReactMarkdown
+                            components={{
+                              p: ({ children }) => <p className="mb-2 text-sm last:mb-0">{children}</p>,
+                              ul: ({ children }) => <ul className="mb-2 list-inside list-disc text-sm">{children}</ul>,
+                              ol: ({ children }) => (
+                                <ol className="mb-2 list-inside list-decimal text-sm">{children}</ol>
+                              ),
+                              li: ({ children }) => <li className="mb-1">{children}</li>,
+                              h1: ({ children }) => <h1 className="mb-2 text-xl font-bold">{children}</h1>,
+                              h2: ({ children }) => <h2 className="mb-2 text-lg font-semibold">{children}</h2>,
+                              h3: ({ children }) => <h3 className="mb-2 text-base font-medium">{children}</h3>,
+                              code: ({ children }) => (
+                                <code className="rounded bg-black px-1 py-0.5 font-mono text-sm text-white">
+                                  {children}
+                                </code>
+                              )
+                            }}
+                          >
+                            {message.content}
+                          </ReactMarkdown>
+                        ) : part.type === 'tool-invocation' ? (
+                          <div className="text-muted-foreground mt-1 text-xs">
+                            <Clock className="mr-1 inline-block h-4 w-4" />
+                            Tool called: {part.toolInvocation.toolName}
+                          </div>
+                        ) : part.type === 'reasoning' ? (
+                          <div className="text-muted-foreground mt-1 text-xs">
+                            <Clock className="mr-1 inline-block h-4 w-4" />
+                            {part.type}
+                          </div>
+                        ) : null}
+                      </div>
+                    </div>
+                  ))}
                 </div>
               ))}
               {isLoading && (