Improve Multi Agent Support

frontend/src/components/Tasklist/Task.tsx

@@ -1,3 +1,5 @@
+import { Markdown } from '@/components/Markdown';
+
 import { TaskStatusIcon } from './TaskStatusIcon';
 
 export interface ITask {
@@ -14,9 +16,11 @@ export interface ITaskList {
 interface TaskProps {
   index: number;
   task: ITask;
+  allowHtml?: boolean;
+  latex?: boolean;
 }
 
-export const Task = ({ index, task }: TaskProps) => {
+export const Task = ({ index, task, allowHtml, latex }: TaskProps) => {
   const statusStyles = {
     ready: '',
     running: 'font-semibold',
@@ -48,14 +52,26 @@ export const Task = ({ index, task }: TaskProps) => {
   return (
     <div className={`task task-status-${task.status}`}>
       <div
-        className={`w-full flex font-medium py-2 text-sm leading-snug ${
+        className={`w-full grid grid-cols-[auto_auto_1fr] items-start gap-1.5 font-medium py-0.5 px-1 text-sm leading-tight ${
           statusStyles[task.status]
         } ${task.forId ? 'cursor-pointer' : 'cursor-default'}`}
         onClick={handleClick}
       >
-        <span className="flex-none w-8 pr-2">{index}</span>
-        <TaskStatusIcon status={task.status} />
-        <span className="pl-2">{task.title}</span>
+        <div className="text-xs text-muted-foreground text-right pr-1 pt-[1px]">
+          {index}
+        </div>
+        <div className="flex items-start pt-[1px]">
+          <TaskStatusIcon status={task.status} />
+        </div>
+        <div className="min-w-0">
+          <Markdown
+            allowHtml={allowHtml}
+            latex={latex}
+            className="max-w-none prose-sm text-left break-words [&_p]:m-0 [&_p]:leading-snug [&_div]:leading-snug [&_div]:mt-0 [&_strong]:font-semibold"
+          >
+            {task.title}
+          </Markdown>
+        </div>
       </div>
     </div>
   );

frontend/src/components/Tasklist/index.tsx

@@ -1,7 +1,7 @@
 import { cn } from '@/lib/utils';
 import useSWR from 'swr';
 
-import { useChatData } from '@chainlit/react-client';
+import { useChatData, useConfig } from '@chainlit/react-client';
 
 import { Badge } from '@/components/ui/badge';
 import { Card, CardContent, CardHeader } from '@/components/ui/card';
@@ -17,7 +17,7 @@ const fetcher = (url: string) =>
 
 const Header = ({ status }: HeaderProps) => {
   return (
-    <CardHeader className="flex flex-row items-center justify-between">
+    <CardHeader className="flex flex-row items-center justify-between gap-2 p-3">
       <div className="font-semibold">Tasks</div>
       <Badge variant="secondary">{status || '?'}</Badge>
     </CardHeader>
@@ -32,6 +32,10 @@ interface TaskListProps {
 const TaskList = ({ isMobile, isCopilot }: TaskListProps) => {
   const { tasklists } = useChatData();
   const tasklist = tasklists[tasklists.length - 1];
+  const { config } = useConfig();
+
+  const allowHtml = config?.features?.unsafe_allow_html;
+  const latex = config?.features?.latex;
 
   const { error, data, isLoading } = useSWR<ITaskList>(tasklist?.url, fetcher, {
     keepPreviousData: true
@@ -70,8 +74,13 @@ const TaskList = ({ isMobile, isCopilot }: TaskListProps) => {
         <Card>
           <Header status={content.status} />
           {highlightedTask && (
-            <CardContent>
-              <Task index={highlightedTaskIndex + 1} task={highlightedTask} />
+            <CardContent className="p-2.5">
+              <Task
+                index={highlightedTaskIndex + 1}
+                task={highlightedTask}
+                allowHtml={allowHtml}
+                latex={latex}
+              />
             </CardContent>
           )}
         </Card>
@@ -80,12 +89,18 @@ const TaskList = ({ isMobile, isCopilot }: TaskListProps) => {
   }
 
   return (
-    <aside className="hidden tasklist max-w-96 flex-grow md:block overflow-y-auto mr-4 mb-4">
+    <aside className="hidden tasklist max-w-[21rem] flex-grow md:block overflow-y-auto mr-3 mb-3">
       <Card className="overflow-y-auto h-full">
         <Header status={content?.status} />
-        <CardContent className="flex flex-col gap-2">
+        <CardContent className="flex flex-col gap-1 p-2.5">
           {tasks?.map((task, index) => (
-            <Task key={index} index={index + 1} task={task} />
+            <Task
+              key={index}
+              index={index + 1}
+              task={task}
+              allowHtml={allowHtml}
+              latex={latex}
+            />
           ))}
         </CardContent>
       </Card>

frontend/src/components/chat/Messages/Message/Content/index.tsx

@@ -1,5 +1,5 @@
 import { prepareContent } from '@/lib/message';
-import { forwardRef, memo } from 'react';
+import { forwardRef, memo, useMemo } from 'react';
 
 import type { IMessageElement, IStep } from '@chainlit/react-client';
 
@@ -8,16 +8,19 @@ import { Markdown } from '@/components/Markdown';
 
 import { InlinedElements } from './InlinedElements';
 
+type ContentSection = 'input' | 'output';
+
 export interface Props {
   elements: IMessageElement[];
   message: IStep;
   allowHtml?: boolean;
   latex?: boolean;
+  sections?: ContentSection[];
 }
 
 const MessageContent = memo(
   forwardRef<HTMLDivElement, Props>(
-    ({ message, elements, allowHtml, latex }, ref) => {
+    ({ message, elements, allowHtml, latex, sections }, ref) => {
       const outputContent =
         message.streaming && message.output
           ? message.output + CURSOR_PLACEHOLDER
@@ -34,11 +37,19 @@ const MessageContent = memo(
         language: message.language
       });
 
-      const displayInput = message.input && message.showInput;
+      const selectedSections = sections ?? ['input', 'output'];
+      const sectionsSet = useMemo(
+        () => new Set(selectedSections),
+        [selectedSections]
+      );
+
+      const displayInput =
+        sectionsSet.has('input') && message.input && message.showInput;
+      const displayOutput = sectionsSet.has('output');
 
       const isMessage = message.type.includes('message');
 
-      const outputMarkdown = (
+      const outputMarkdown = displayOutput ? (
         <>
           {!isMessage && displayInput && message.output ? (
             <div className="font-medium">Output</div>
@@ -51,7 +62,7 @@ const MessageContent = memo(
             {output}
           </Markdown>
         </>
-      );
+      ) : null;
 
       let inputMarkdown;
 
@@ -73,8 +84,6 @@ const MessageContent = memo(
 
         inputMarkdown = (
           <>
-            <div className="font-medium">Input</div>
-
             <Markdown
               allowHtml={allowHtml}
               latex={latex}
@@ -95,8 +104,10 @@ const MessageContent = memo(
 
       return (
         <div ref={ref} className="message-content w-full flex flex-col gap-2">
-          {!!inputMarkdown || output ? markdownContent : null}
-          <InlinedElements elements={outputInlinedElements} />
+          {displayInput || (displayOutput && output) ? markdownContent : null}
+          {displayOutput ? (
+            <InlinedElements elements={outputInlinedElements} />
+          ) : null}
         </div>
       );
     }

frontend/src/components/chat/Messages/Message/index.tsx

@@ -51,6 +51,8 @@ const Message = memo(
     const hiddenSkip = isStep && cot === 'hidden';
 
     const skip = toolCallSkip || hiddenSkip;
+    const showInputSection = Boolean(message.input && message.showInput);
+    const shouldRenderOutput = !showInputSection || Boolean(message.output);
 
     if (skip) {
       if (!message.steps) {
@@ -104,6 +106,15 @@ const Message = memo(
                   {/* Display the step and its children */}
                   {isStep ? (
                     <Step step={message} isRunning={isRunning}>
+                      {showInputSection ? (
+                        <MessageContent
+                          elements={elements}
+                          message={message}
+                          allowHtml={allowHtml}
+                          latex={latex}
+                          sections={['input']}
+                        />
+                      ) : null}
                       {message.steps ? (
                         <Messages
                           messages={message.steps.filter(
@@ -115,13 +126,16 @@ const Message = memo(
                           isRunning={isRunning}
                         />
                       ) : null}
-                      <MessageContent
-                        ref={contentRef}
-                        elements={elements}
-                        message={message}
-                        allowHtml={allowHtml}
-                        latex={latex}
-                      />
+                      {shouldRenderOutput ? (
+                        <MessageContent
+                          ref={contentRef}
+                          elements={elements}
+                          message={message}
+                          allowHtml={allowHtml}
+                          latex={latex}
+                          sections={showInputSection ? ['output'] : undefined}
+                        />
+                      ) : null}
                       <MessageButtons
                         message={message}
                         actions={actions}