feat(web): enhance WebSocket handling and chat panel improvements

- Add react-markdown with remark-breaks and remark-gfm for markdown rendering - Improve WebSocket client with better reconnection handling - Enhanced task chat panel with markdown support for comments - Add comment type and metadata support for tool call rendering - Improve connection status indicator - Update task page with improved agent controls - Add WebSocket router for better message handling - Update state store with enhanced type safety
2026-01-12 02:05:30 +00:00
parent 00c0a0c529
commit 5381094bba
18 changed files with 2280 additions and 208 deletions
@@ -2,6 +2,7 @@ import { PageClient } from '@/app/page-client';
 import { StateHydrator } from '@/components/state-hydrator';
 import { fetchBoardSnapshot, fetchBoards, fetchWorkspaces } from '@/lib/ssr/http';
 import { snapshotToState } from '@/lib/ssr/mapper';
+import type { AppState } from '@/lib/state/store';

 // Server Component: runs on the server for SSR and data hydration.
 type PageProps = {
@@ -21,7 +22,7 @@ export default async function Page({ searchParams }: PageProps) {
      workspaces.workspaces.find((workspace) => workspace.id === workspaceId)?.id ??
      workspaces.workspaces[0]?.id ??
      null;
-    let initialState = {
+    let initialState: Partial<AppState> = {
      workspaces: {
        items: workspaces.workspaces.map((workspace) => ({
          id: workspace.id,
@@ -627,7 +627,7 @@ export function WorkspaceEditClient({ workspace, boards, repositories }: Workspa
  };

  const handleDeleteWorkspace = async () => {
-    if (deleteConfirmText !== 'delete') return;
+    if (deleteConfirmText !== 'delete' || !currentWorkspace) return;
    try {
      await deleteWorkspaceRequest.run(currentWorkspace.id);
      setWorkspaces(workspaces.filter((workspace) => workspace.id !== currentWorkspace.id));
@@ -9,79 +9,51 @@ import { Textarea } from '@/components/ui/textarea';
 import { getLocalStorage, setLocalStorage } from '@/lib/local-storage';
 import { TaskChatPanel } from '@/components/task/task-chat-panel';
 import { TaskTopBar } from '@/components/task/task-top-bar';
-import type { Task } from '@/lib/types/http';
+import type { Task, Comment } from '@/lib/types/http';
 import { TaskFilesPanel } from '@/components/task/task-files-panel';
 import { TaskChangesPanel } from '@/components/task/task-changes-panel';
 import { TaskRightPanel } from '@/components/task/task-right-panel';
 import { getBackendConfig } from '@/lib/config';
 import { listRepositories, listRepositoryBranches } from '@/lib/http/client';
 import { useRequest } from '@/lib/http/use-request';
-
-type ChatMessage = {
-  id: string;
-  role: 'user' | 'agent';
-  content: string;
-};
-
-type ChatSession = {
-  id: string;
-  title: string;
-  messages: ChatMessage[];
-};
+import { getWebSocketClient } from '@/lib/ws/connection';
+import { useAppStoreApi } from '@/components/state-provider';
+import { useAppStore } from '@/components/state-provider';

 const AGENTS = [
  { id: 'codex', label: 'Codex' },
  { id: 'claude', label: 'Claude Code' },
 ];

-const INITIAL_CHATS: ChatSession[] = [
-  {
-    id: 'chat-1',
-    title: 'Build overview',
-    messages: [
-      {
-        id: 'm1',
-        role: 'agent',
-        content: 'I can help with this task. What should I tackle first?',
-      },
-    ],
-  },
-];
-
 type TaskPageClientProps = {
  task: Task | null;
 };

-function buildInitialChats(taskPrompt?: string): ChatSession[] {
-  if (!taskPrompt) {
-    return INITIAL_CHATS;
-  }
-  return [
-    {
-      ...INITIAL_CHATS[0],
-      messages: [
-        {
-          id: 'prompt',
-          role: 'user',
-          content: taskPrompt,
-        },
-        ...INITIAL_CHATS[0].messages,
-      ],
-    },
-  ];
-}
-
 export default function TaskPage({ task }: TaskPageClientProps) {
+  const store = useAppStoreApi();
  const defaultHorizontalLayout: [number, number] = [75, 25];
  const [horizontalLayout, setHorizontalLayout] = useState<[number, number]>(() =>
    getLocalStorage('task-layout-horizontal', defaultHorizontalLayout)
  );
-  const [chats, setChats] = useState<ChatSession[]>(() =>
-    buildInitialChats(task?.description ?? '')
-  );
  const [leftTab, setLeftTab] = useState('chat');
  const [selectedDiffPath, setSelectedDiffPath] = useState<string | null>(null);
  const [notes, setNotes] = useState('');
+  const [isLoadingComments, setIsLoadingComments] = useState(false);
+  const [isAgentRunning, setIsAgentRunning] = useState(false);
+  const [isAgentLoading, setIsAgentLoading] = useState(false);
+
+  // Track task state from store (kanban.tasks) to determine if agent is running
+  const taskFromStore = useAppStore((state) =>
+    state.kanban.tasks.find((t) => t.id === task?.id)
+  );
+
+  // Update isAgentRunning based on task state
+  useEffect(() => {
+    // Check if task state indicates agent is working
+    const taskState = task?.state;
+    setIsAgentRunning(taskState === 'IN_PROGRESS');
+  }, [task?.state]);
+
  const fetchBranches = useCallback(async (workspaceId: string, repoPath: string) => {
    const response = await listRepositories(getBackendConfig().apiBaseUrl, workspaceId);
    const repo = response.repositories.find((item) => item.local_path === repoPath);
@@ -91,27 +63,105 @@ export default function TaskPage({ task }: TaskPageClientProps) {
  }, []);
  const branchesRequest = useRequest(fetchBranches);

-  const activeChat = chats[0];
+  // Fetch comments on mount and when task changes
+  useEffect(() => {
+    if (!task?.id) return;

-  const handleSendMessage = useCallback((content: string) => {
-    setChats((currentChats) => {
-      const [first, ...rest] = currentChats;
-      if (!first) return currentChats;
-      return [
-        {
-          ...first,
-          messages: [...first.messages, { id: crypto.randomUUID(), role: 'user', content }],
-        },
-        ...rest,
-      ];
-    });
-  }, []);
+    const fetchComments = async () => {
+      const client = getWebSocketClient();
+      if (!client) return;
+
+      setIsLoadingComments(true);
+      store.getState().setCommentsLoading(true);
+
+      try {
+        const response = await client.request<Comment[]>('comment.list', { task_id: task.id }, 10000);
+        console.log('[API] comment.list response:', JSON.stringify(response, null, 2));
+        store.getState().setComments(task.id, response);
+      } catch (error) {
+        console.error('Failed to fetch comments:', error);
+        store.getState().setComments(task.id, []);
+      } finally {
+        setIsLoadingComments(false);
+      }
+    };
+
+    fetchComments();
+  }, [task?.id, store]);
+
+  // Subscribe to task for real-time updates
+  useEffect(() => {
+    if (!task?.id) return;
+
+    const client = getWebSocketClient();
+    if (!client) return;
+
+    // Subscribe to task updates
+    client.subscribe(task.id);
+
+    return () => {
+      // Unsubscribe when leaving
+      client.unsubscribe(task.id);
+    };
+  }, [task?.id]);
+
+  const handleSendMessage = useCallback(async (content: string) => {
+    if (!task?.id) return;
+
+    const client = getWebSocketClient();
+    if (!client) return;
+
+    try {
+      await client.request('comment.add', { task_id: task.id, content }, 10000);
+    } catch (error) {
+      console.error('Failed to send comment:', error);
+    }
+  }, [task?.id]);

  const handleSelectDiffPath = useCallback((path: string) => {
    setSelectedDiffPath(path);
    setLeftTab('changes');
  }, []);

+  const handleStartAgent = useCallback(async () => {
+    if (!task?.id) return;
+
+    const client = getWebSocketClient();
+    if (!client) return;
+
+    setIsAgentLoading(true);
+    try {
+      // Use task's agent_type if set, otherwise default to 'augment-agent'
+      const agentType = task.agent_type ?? 'augment-agent';
+      await client.request('orchestrator.start', {
+        task_id: task.id,
+        agent_type: agentType,
+      }, 15000);
+      setIsAgentRunning(true);
+    } catch (error) {
+      console.error('Failed to start agent:', error);
+    } finally {
+      setIsAgentLoading(false);
+    }
+  }, [task?.id, task?.agent_type]);
+
+  const handleStopAgent = useCallback(async () => {
+    if (!task?.id) return;
+
+    const client = getWebSocketClient();
+    if (!client) return;
+
+    setIsAgentLoading(true);
+    try {
+      await client.request('orchestrator.stop', { task_id: task.id }, 15000);
+      setIsAgentRunning(false);
+    } catch (error) {
+      console.error('Failed to stop agent:', error);
+    } finally {
+      setIsAgentLoading(false);
+    }
+  }, [task?.id]);
+
  const topFilesPanel = <TaskFilesPanel onSelectDiffPath={handleSelectDiffPath} />;

  useEffect(() => {
@@ -127,6 +177,10 @@ export default function TaskPage({ task }: TaskPageClientProps) {
        baseBranch={task?.branch ?? undefined}
        branches={task?.repository_url ? branchesRequest.data ?? [] : []}
        branchesLoading={branchesRequest.isLoading}
+        onStartAgent={handleStartAgent}
+        onStopAgent={handleStopAgent}
+        isAgentRunning={isAgentRunning}
+        isAgentLoading={isAgentLoading}
      />

      <div className="flex-1 min-h-0 px-4 pb-4">
@@ -153,7 +207,7 @@ export default function TaskPage({ task }: TaskPageClientProps) {
                    All changes
                  </TabsTrigger>
                  <TabsTrigger value="chat" className="cursor-pointer">
-                    Chat {activeChat ? `• ${activeChat.title}` : ''}
+                    Chat
                  </TabsTrigger>
                </TabsList>

@@ -174,7 +228,18 @@ export default function TaskPage({ task }: TaskPageClientProps) {
                </TabsContent>

                <TabsContent value="chat" className="mt-3 flex flex-col min-h-0 flex-1">
-                  <TaskChatPanel activeChat={activeChat} agents={AGENTS} onSend={handleSendMessage} />
+                  {task?.id ? (
+                    <TaskChatPanel
+                      taskId={task.id}
+                      agents={AGENTS}
+                      onSend={handleSendMessage}
+                      isLoading={isLoadingComments}
+                    />
+                  ) : (
+                    <div className="flex items-center justify-center h-full text-muted-foreground">
+                      No task selected
+                    </div>
+                  )}
                </TabsContent>
              </Tabs>
            </div>
@@ -14,6 +14,8 @@ export function ConnectionStatus() {
        return { color: 'text-emerald-500', text: 'Connected', variant: 'default' as const, pulse: false };
      case 'connecting':
        return { color: 'text-amber-500', text: 'Connecting', variant: 'secondary' as const, pulse: true };
+      case 'reconnecting':
+        return { color: 'text-orange-500', text: 'Reconnecting', variant: 'secondary' as const, pulse: true };
      case 'error':
        return { color: 'text-red-500', text: 'Error', variant: 'destructive' as const, pulse: false };
      case 'disconnected':
@@ -48,7 +48,7 @@ export function RepositoryCard({
  const { toast } = useToast();
  const [deleteOpen, setDeleteOpen] = useState(false);
  const saveRequest = useRequest(() => onSave(repository.id));
-  const deleteRequest = useRequest(() => onDelete(repository.id));
+  const deleteRequest = useRequest(async () => { await onDelete(repository.id); });

  const handleSave = async () => {
    try {
@@ -9,30 +9,33 @@ import {
 } from '@/components/ui/breadcrumb';
 import { SidebarInset, SidebarProvider } from '@/components/ui/sidebar';
 import { SettingsAppSidebar } from '@/components/settings/settings-app-sidebar';
+import { TooltipProvider } from '@/components/ui/tooltip';

 export function SettingsLayoutClient({ children }: { children: React.ReactNode }) {
  return (
-    <SidebarProvider>
-      <SettingsAppSidebar />
-      <SidebarInset>
-        <header className="flex h-16 shrink-0 items-center gap-2">
-          <div className="flex items-center gap-2 px-4">
-            <Breadcrumb>
-              <BreadcrumbList>
-                <BreadcrumbItem>
-                  <BreadcrumbLink href="/" className="flex items-center gap-2">
-                    <IconArrowLeft className="h-4 w-4" />
-                    Home
-                  </BreadcrumbLink>
-                </BreadcrumbItem>
-              </BreadcrumbList>
-            </Breadcrumb>
+    <TooltipProvider>
+      <SidebarProvider>
+        <SettingsAppSidebar />
+        <SidebarInset>
+          <header className="flex h-16 shrink-0 items-center gap-2">
+            <div className="flex items-center gap-2 px-4">
+              <Breadcrumb>
+                <BreadcrumbList>
+                  <BreadcrumbItem>
+                    <BreadcrumbLink href="/" className="flex items-center gap-2">
+                      <IconArrowLeft className="h-4 w-4" />
+                      Home
+                    </BreadcrumbLink>
+                  </BreadcrumbItem>
+                </BreadcrumbList>
+              </Breadcrumb>
+            </div>
+          </header>
+          <div className="flex flex-1 flex-col gap-4 p-4 pt-0">
+            {children}
          </div>
-        </header>
-        <div className="flex flex-1 flex-col gap-4 p-4 pt-0">
-          {children}
-        </div>
-      </SidebarInset>
-    </SidebarProvider>
+        </SidebarInset>
+      </SidebarProvider>
+    </TooltipProvider>
  );
 }
@@ -1,7 +1,10 @@
 'use client';

-import { FormEvent, useState } from 'react';
-import { IconBrain, IconListCheck, IconPaperclip } from '@tabler/icons-react';
+import { FormEvent, useEffect, useRef, useState } from 'react';
+import { IconBrain, IconListCheck, IconLoader2, IconPaperclip, IconTerminal2 } from '@tabler/icons-react';
+import ReactMarkdown from 'react-markdown';
+import remarkBreaks from 'remark-breaks';
+import remarkGfm from 'remark-gfm';
 import { Button } from '@/components/ui/button';
 import { Textarea } from '@/components/ui/textarea';
 import {
@@ -13,18 +16,8 @@ import {
 import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
 import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip';
 import { cn } from '@/lib/utils';
-
-type ChatMessage = {
-  id: string;
-  role: 'user' | 'agent';
-  content: string;
-};
-
-type ChatSession = {
-  id: string;
-  title: string;
-  messages: ChatMessage[];
-};
+import { useAppStore } from '@/components/state-provider';
+import type { Comment } from '@/lib/types/http';

 type AgentOption = {
  id: string;
@@ -32,43 +25,111 @@ type AgentOption = {
 };

 type TaskChatPanelProps = {
-  activeChat?: ChatSession;
+  taskId: string;
  agents: AgentOption[];
  onSend: (message: string) => void;
+  isLoading?: boolean;
 };

-export function TaskChatPanel({ activeChat, agents, onSend }: TaskChatPanelProps) {
+export function TaskChatPanel({ taskId, agents, onSend, isLoading }: TaskChatPanelProps) {
  const [messageInput, setMessageInput] = useState('');
  const [selectedAgent, setSelectedAgent] = useState(agents[0]?.id ?? '');
  const [planModeEnabled, setPlanModeEnabled] = useState(false);
+  const [isSending, setIsSending] = useState(false);
+  const messagesEndRef = useRef<HTMLDivElement>(null);

-  const handleSubmit = (event: FormEvent) => {
+  const commentsState = useAppStore((state) => state.comments);
+  const comments = commentsState?.items ?? [];
+  const commentsLoading = commentsState?.isLoading ?? false;
+
+  // Filter to only show message, content, and tool_call types (not progress, etc)
+  const visibleComments = comments.filter(
+    (c) => c.type === 'message' || c.type === 'content' || c.type === 'tool_call' || !c.type
+  );
+
+  // Scroll to bottom when new comments arrive
+  useEffect(() => {
+    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
+  }, [visibleComments.length]);
+
+  const handleSubmit = async (event: FormEvent) => {
    event.preventDefault();
    const trimmed = messageInput.trim();
-    if (!trimmed) return;
-    onSend(trimmed);
+    if (!trimmed || isSending) return;
+    setIsSending(true);
    setMessageInput('');
+    try {
+      await onSend(trimmed);
+    } finally {
+      setIsSending(false);
+    }
+  };
+
+  const formatContent = (comment: Comment) => {
+    // For content type comments, the content might be streamed chunks
+    // For message type, show as-is
+    return comment.content || '(empty)';
  };

  return (
    <>
      <div className="flex-1 min-h-0 overflow-y-auto rounded-lg bg-background p-3 space-y-3">
-        {(activeChat?.messages ?? []).map((message) => (
-          <div
-            key={message.id}
-            className={cn(
-              'max-w-[80%] rounded-lg px-3 py-2 text-sm leading-relaxed',
-              message.role === 'user'
-                ? 'ml-auto bg-primary text-primary-foreground'
-                : 'bg-muted text-foreground'
-            )}
-          >
-            <p className="text-[11px] uppercase tracking-wide mb-1 opacity-70">
-              {message.role === 'user' ? 'You' : 'Agent'}
-            </p>
-            <p>{message.content}</p>
+        {(isLoading || commentsLoading) && (
+          <div className="flex items-center justify-center py-8 text-muted-foreground">
+            <IconLoader2 className="h-5 w-5 animate-spin mr-2" />
+            <span>Loading comments...</span>
          </div>
-        ))}
+        )}
+        {!isLoading && !commentsLoading && visibleComments.length === 0 && (
+          <div className="flex items-center justify-center py-8 text-muted-foreground">
+            <span>No messages yet. Start the conversation!</span>
+          </div>
+        )}
+        {visibleComments.map((comment) => {
+          // Tool call comments have a special rendering
+          if (comment.type === 'tool_call') {
+            const metadata = comment.metadata as { title?: string; status?: string } | undefined;
+            return (
+              <div
+                key={comment.id}
+                className="flex items-center gap-2 px-3 py-2 text-sm rounded-md bg-muted/30 border border-border/30"
+              >
+                <IconTerminal2 className="h-4 w-4 text-amber-600 dark:text-amber-400 flex-shrink-0" />
+                <span className="font-mono text-xs text-muted-foreground">
+                  {metadata?.title || comment.content}
+                </span>
+              </div>
+            );
+          }
+
+          // Regular message or agent response - render as markdown
+          const isUser = comment.author_type === 'user';
+          return (
+            <div
+              key={comment.id}
+              className={cn(
+                'max-w-[85%] rounded-lg px-4 py-3 text-sm',
+                isUser
+                  ? 'ml-auto bg-primary text-primary-foreground'
+                  : 'bg-muted text-foreground'
+              )}
+            >
+              <p className="text-[11px] uppercase tracking-wide mb-2 opacity-70">
+                {isUser ? 'You' : 'Agent'}
+              </p>
+              {isUser ? (
+                <p className="whitespace-pre-wrap">{formatContent(comment)}</p>
+              ) : (
+                <div className="prose prose-sm dark:prose-invert max-w-none prose-p:my-2 prose-p:leading-relaxed prose-ul:my-2 prose-ol:my-2 prose-li:my-0.5 prose-pre:my-3 prose-code:px-1.5 prose-code:py-0.5 prose-code:bg-background/50 prose-code:rounded prose-code:text-xs prose-code:before:content-none prose-code:after:content-none prose-pre:bg-background/80 prose-pre:text-xs prose-strong:text-foreground prose-headings:text-foreground">
+                  <ReactMarkdown remarkPlugins={[remarkGfm, remarkBreaks]}>
+                    {formatContent(comment)}
+                  </ReactMarkdown>
+                </div>
+              )}
+            </div>
+          );
+        })}
+        <div ref={messagesEndRef} />
      </div>
      <form onSubmit={handleSubmit} className="mt-3 flex flex-col gap-2">
        <Textarea
@@ -14,7 +14,10 @@ import {
  IconGitFork,
  IconGitMerge,
  IconGitPullRequest,
+  IconLoader2,
  IconPencil,
+  IconPlayerPlay,
+  IconPlayerStop,
 } from '@tabler/icons-react';
 import { Button } from '@/components/ui/button';
 import { CommitStatBadge, LineStat } from '@/components/diff-stat';
@@ -27,6 +30,10 @@ type TaskTopBarProps = {
  baseBranch?: string;
  branches?: string[];
  branchesLoading?: boolean;
+  onStartAgent?: () => void;
+  onStopAgent?: () => void;
+  isAgentRunning?: boolean;
+  isAgentLoading?: boolean;
 };

 const TaskTopBar = memo(function TaskTopBar({
@@ -34,6 +41,10 @@ const TaskTopBar = memo(function TaskTopBar({
  baseBranch,
  branches = [],
  branchesLoading = false,
+  onStartAgent,
+  onStopAgent,
+  isAgentRunning = false,
+  isAgentLoading = false,
 }: TaskTopBarProps) {
  const [branchName, setBranchName] = useState('feature/agent-ui');
  const [isEditingBranch, setIsEditingBranch] = useState(false);
@@ -156,6 +167,48 @@ const TaskTopBar = memo(function TaskTopBar({
        </Tooltip>
      </div>
      <div className="flex items-center gap-2">
+        {/* Start/Stop Agent Button */}
+        {isAgentRunning ? (
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                size="sm"
+                variant="destructive"
+                className="cursor-pointer"
+                onClick={onStopAgent}
+                disabled={isAgentLoading}
+              >
+                {isAgentLoading ? (
+                  <IconLoader2 className="h-4 w-4 animate-spin" />
+                ) : (
+                  <IconPlayerStop className="h-4 w-4" />
+                )}
+                Stop Agent
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Stop the agent</TooltipContent>
+          </Tooltip>
+        ) : (
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                size="sm"
+                variant="default"
+                className="cursor-pointer"
+                onClick={onStartAgent}
+                disabled={isAgentLoading}
+              >
+                {isAgentLoading ? (
+                  <IconLoader2 className="h-4 w-4 animate-spin" />
+                ) : (
+                  <IconPlayerPlay className="h-4 w-4" />
+                )}
+                Start Agent
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Start the agent on this task</TooltipContent>
+          </Tooltip>
+        )}
        <Tooltip>
          <TooltipTrigger asChild>
            <Button size="sm" variant="outline" className="cursor-pointer">
@@ -1,6 +1,8 @@
-import type { AppState } from '@/lib/state/store';
+import type { AppState, KanbanState } from '@/lib/state/store';
 import type { BoardSnapshot, Task } from '@/lib/types/http';

+type KanbanTask = KanbanState['tasks'][number];
+
 export function snapshotToState(snapshot: BoardSnapshot): Partial<AppState> {
  const tasks = snapshot.tasks
    .map((task) => {
@@ -13,20 +15,9 @@ export function snapshotToState(snapshot: BoardSnapshot): Partial<AppState> {
        description: task.description ?? undefined,
        position: task.position ?? 0,
        state: task.state,
-      };
+      } as KanbanTask;
    })
-    .filter(
-      (
-        task
-      ): task is {
-        id: string;
-        columnId: string;
-        title: string;
-        description?: string;
-        position: number;
-        state?: Task['state'];
-      } => Boolean(task)
-    );
+    .filter((task): task is KanbanTask => task !== null);

  return {
    kanban: {
@@ -1,5 +1,5 @@
 import type { ReactNode } from 'react';
-import type { TaskState as TaskStatus } from '@/lib/types/http';
+import type { TaskState as TaskStatus, Comment } from '@/lib/types/http';
 import { createStore } from 'zustand/vanilla';
 import { immer } from 'zustand/middleware/immer';

@@ -43,10 +43,16 @@ export type DiffState = {
 };

 export type ConnectionState = {
-  status: 'disconnected' | 'connecting' | 'connected' | 'error';
+  status: 'disconnected' | 'connecting' | 'connected' | 'error' | 'reconnecting';
  error: string | null;
 };

+export type CommentsState = {
+  taskId: string | null;
+  items: Comment[];
+  isLoading: boolean;
+};
+
 export type AppState = {
  kanban: KanbanState;
  workspaces: WorkspaceState;
@@ -56,6 +62,7 @@ export type AppState = {
  terminal: TerminalState;
  diffs: DiffState;
  connection: ConnectionState;
+  comments: CommentsState;
  hydrate: (state: Partial<AppState>) => void;
  setActiveWorkspace: (workspaceId: string | null) => void;
  setWorkspaces: (workspaces: WorkspaceState['items']) => void;
@@ -63,6 +70,9 @@ export type AppState = {
  setBoards: (boards: BoardState['items']) => void;
  setTerminalOutput: (terminalId: string, data: string) => void;
  setConnectionStatus: (status: ConnectionState['status'], error?: string | null) => void;
+  setComments: (taskId: string, comments: Comment[]) => void;
+  addComment: (comment: Comment) => void;
+  setCommentsLoading: (loading: boolean) => void;
 };

 export type AppStore = ReturnType<typeof createAppStore>;
@@ -76,6 +86,7 @@ const defaultState: AppState = {
  terminal: { terminals: [] },
  diffs: { files: [] },
  connection: { status: 'disconnected', error: null },
+  comments: { taskId: null, items: [], isLoading: false },
  hydrate: () => undefined,
  setActiveWorkspace: () => undefined,
  setWorkspaces: () => undefined,
@@ -83,16 +94,43 @@ const defaultState: AppState = {
  setBoards: () => undefined,
  setTerminalOutput: () => undefined,
  setConnectionStatus: () => undefined,
+  setComments: () => undefined,
+  addComment: () => undefined,
+  setCommentsLoading: () => undefined,
 };

+function mergeInitialState(initialState?: Partial<AppState>): Omit<AppState, 'hydrate' | 'setActiveWorkspace' | 'setWorkspaces' | 'setActiveBoard' | 'setBoards' | 'setTerminalOutput' | 'setConnectionStatus' | 'setComments' | 'addComment' | 'setCommentsLoading'> {
+  if (!initialState) return defaultState;
+  return {
+    workspaces: { ...defaultState.workspaces, ...initialState.workspaces },
+    boards: { ...defaultState.boards, ...initialState.boards },
+    kanban: { ...defaultState.kanban, ...initialState.kanban },
+    tasks: { ...defaultState.tasks, ...initialState.tasks },
+    agents: { ...defaultState.agents, ...initialState.agents },
+    terminal: { ...defaultState.terminal, ...initialState.terminal },
+    diffs: { ...defaultState.diffs, ...initialState.diffs },
+    connection: { ...defaultState.connection, ...initialState.connection },
+    comments: { ...defaultState.comments, ...initialState.comments },
+  };
+}
+
 export function createAppStore(initialState?: Partial<AppState>) {
+  const merged = mergeInitialState(initialState);
  return createStore<AppState>()(
    immer((set) => ({
-      ...defaultState,
-      ...initialState,
+      ...merged,
      hydrate: (state) =>
        set((draft) => {
-          Object.assign(draft, state);
+          // Deep merge to avoid overwriting nested objects with undefined
+          if (state.workspaces) Object.assign(draft.workspaces, state.workspaces);
+          if (state.boards) Object.assign(draft.boards, state.boards);
+          if (state.kanban) Object.assign(draft.kanban, state.kanban);
+          if (state.tasks) Object.assign(draft.tasks, state.tasks);
+          if (state.agents) Object.assign(draft.agents, state.agents);
+          if (state.terminal) Object.assign(draft.terminal, state.terminal);
+          if (state.diffs) Object.assign(draft.diffs, state.diffs);
+          if (state.connection) Object.assign(draft.connection, state.connection);
+          if (state.comments) Object.assign(draft.comments, state.comments);
        }),
      setActiveWorkspace: (workspaceId) =>
        set((draft) => {
@@ -130,6 +168,27 @@ export function createAppStore(initialState?: Partial<AppState>) {
          draft.connection.status = status;
          draft.connection.error = error;
        }),
+      setComments: (taskId, comments) =>
+        set((draft) => {
+          draft.comments.taskId = taskId;
+          draft.comments.items = comments;
+          draft.comments.isLoading = false;
+        }),
+      addComment: (comment) =>
+        set((draft) => {
+          // Only add if this comment is for the current task
+          if (draft.comments.taskId === comment.task_id) {
+            // Check if comment already exists (avoid duplicates)
+            const exists = draft.comments.items.some((c) => c.id === comment.id);
+            if (!exists) {
+              draft.comments.items.push(comment);
+            }
+          }
+        }),
+      setCommentsLoading: (loading) =>
+        set((draft) => {
+          draft.comments.isLoading = loading;
+        }),
    }))
  );
 }
@@ -16,7 +16,8 @@ export type BackendMessageType =
  | 'board.deleted'
  | 'column.created'
  | 'column.updated'
-  | 'column.deleted';
+  | 'column.deleted'
+  | 'comment.added';

 export type BackendMessage<T extends BackendMessageType, P> = {
  id?: string;
@@ -26,6 +27,8 @@ export type BackendMessage<T extends BackendMessageType, P> = {
  timestamp?: string;
 };

+import type { TaskState } from '@/lib/types/http';
+
 export type KanbanUpdatePayload = {
  boardId: string;
  columns: Array<{ id: string; title: string; color?: string; position?: number }>;
@@ -35,7 +38,7 @@ export type KanbanUpdatePayload = {
    title: string;
    position?: number;
    description?: string;
-    state?: string;
+    state?: TaskState;
  }>;
 };

@@ -45,7 +48,7 @@ export type TaskEventPayload = {
  column_id: string;
  title: string;
  description?: string;
-  state?: string;
+  state?: TaskState;
  priority?: number;
  position?: number;
 };
@@ -106,6 +109,18 @@ export type ColumnPayload = {
  updated_at?: string;
 };

+export type CommentAddedPayload = {
+  task_id: string;
+  comment_id: string;
+  author_type: 'user' | 'agent';
+  author_id?: string;
+  content: string;
+  type?: string;
+  metadata?: Record<string, unknown>;
+  requests_input?: boolean;
+  created_at: string;
+};
+
 export type BackendMessageMap = {
  'kanban.update': BackendMessage<'kanban.update', KanbanUpdatePayload>;
  'task.created': BackendMessage<'task.created', TaskEventPayload>;
@@ -125,4 +140,5 @@ export type BackendMessageMap = {
  'column.created': BackendMessage<'column.created', ColumnPayload>;
  'column.updated': BackendMessage<'column.updated', ColumnPayload>;
  'column.deleted': BackendMessage<'column.deleted', ColumnPayload>;
+  'comment.added': BackendMessage<'comment.added', CommentAddedPayload>;
 };
@@ -141,3 +141,18 @@ export type ListWorkspacesResponse = {
  workspaces: Workspace[];
  total: number;
 };
+
+export type CommentAuthorType = 'user' | 'agent';
+export type CommentType = 'message' | 'content' | 'tool_call' | 'progress' | 'error' | 'status';
+
+export type Comment = {
+  id: string;
+  task_id: string;
+  author_type: CommentAuthorType;
+  author_id?: string;
+  content: string;
+  type: CommentType;
+  metadata?: Record<string, unknown>;
+  requests_input?: boolean;
+  created_at: string;
+};
@@ -0,0 +1,195 @@
+# WebSocket Client with Auto-Reconnect
+
+This directory contains the WebSocket client implementation with automatic reconnection capabilities.
+
+## Features
+
+### 1. Automatic Reconnection
+- **Exponential Backoff**: Reconnection delays increase exponentially to avoid overwhelming the server
+- **Configurable Retry Limits**: Set maximum number of reconnection attempts
+- **Graceful Degradation**: Properly handles connection failures and notifies the UI
+
+### 2. Connection State Management
+The client tracks the following connection states:
+- `idle`: Initial state before any connection attempt
+- `connecting`: Actively establishing a connection
+- `open`: Successfully connected and ready to send/receive messages
+- `reconnecting`: Attempting to reconnect after a disconnect
+- `closed`: Connection closed (intentionally or after max retries)
+- `error`: Connection failed and won't retry
+
+### 3. Message Queue
+- Messages sent while disconnected are queued
+- Queued messages are automatically sent when connection is restored
+- Prevents message loss during brief disconnections
+
+### 4. Subscription Management
+- Tracks task subscriptions across reconnections
+- Automatically re-subscribes to all tasks when reconnected
+- Provides `subscribe()` and `unsubscribe()` methods for easy subscription management
+
+### 5. Pending Request Handling
+- Properly cleans up pending requests on disconnect
+- Rejects pending promises with appropriate error messages
+- Prevents memory leaks from abandoned requests
+
+## Configuration
+
+### Reconnect Options
+
+```typescript
+interface ReconnectOptions {
+  enabled?: boolean;           // Enable/disable auto-reconnect (default: true)
+  maxAttempts?: number;        // Maximum reconnection attempts (default: 10)
+  initialDelay?: number;       // Initial delay in ms (default: 1000)
+  maxDelay?: number;           // Maximum delay in ms (default: 30000)
+  backoffMultiplier?: number;  // Exponential backoff multiplier (default: 1.5)
+}
+```
+
+### Default Configuration
+
+```typescript
+const DEFAULT_RECONNECT_OPTIONS = {
+  enabled: true,
+  maxAttempts: 10,
+  initialDelay: 1000,      // 1 second
+  maxDelay: 30000,         // 30 seconds
+  backoffMultiplier: 1.5,
+};
+```
+
+## Usage
+
+### Basic Usage
+
+```typescript
+import { WebSocketClient } from '@/lib/ws/client';
+
+const client = new WebSocketClient(
+  'ws://localhost:8080/ws',
+  (status) => {
+    console.log('Connection status:', status);
+  }
+);
+
+client.connect();
+```
+
+### Custom Reconnect Configuration
+
+```typescript
+const client = new WebSocketClient(
+  'ws://localhost:8080/ws',
+  (status) => {
+    console.log('Connection status:', status);
+  },
+  {
+    enabled: true,
+    maxAttempts: 5,
+    initialDelay: 2000,
+    maxDelay: 60000,
+    backoffMultiplier: 2,
+  }
+);
+```
+
+### Subscribing to Tasks
+
+```typescript
+// Subscribe to task updates
+client.subscribe('task-id-123');
+
+// Unsubscribe when done
+client.unsubscribe('task-id-123');
+```
+
+### Sending Messages
+
+```typescript
+// Send a message (queued if not connected)
+client.send({ type: 'request', action: 'some.action', payload: {} });
+
+// Send a request and wait for response
+const response = await client.request('task.get', { id: 'task-123' });
+```
+
+### Listening to Events
+
+```typescript
+const unsubscribe = client.on('task.updated', (message) => {
+  console.log('Task updated:', message.payload);
+});
+
+// Clean up when done
+unsubscribe();
+```
+
+## Reconnection Behavior
+
+### Exponential Backoff Example
+
+With default settings:
+- Attempt 1: 1 second delay
+- Attempt 2: 1.5 seconds delay
+- Attempt 3: 2.25 seconds delay
+- Attempt 4: 3.375 seconds delay
+- ...
+- Attempt 10: 30 seconds delay (capped at maxDelay)
+
+### What Happens on Disconnect
+
+1. Connection closes (network issue, server restart, etc.)
+2. Client sets status to `closed`
+3. If not an intentional disconnect and reconnect is enabled:
+   - Calculate delay using exponential backoff
+   - Set status to `reconnecting`
+   - Schedule reconnection attempt
+4. On successful reconnection:
+   - Reset reconnect attempt counter
+   - Set status to `open`
+   - Flush queued messages
+   - Re-subscribe to all tasks
+5. If max attempts reached:
+   - Set status to `error`
+   - Clean up pending requests
+   - Stop reconnection attempts
+
+## Integration with React
+
+The `useWebSocket` hook automatically sets up the client with reconnection:
+
+```typescript
+import { useWebSocket } from '@/lib/ws/use-websocket';
+
+export function MyComponent() {
+  const store = useAppStoreApi();
+  useWebSocket(store, 'ws://localhost:8080/ws');
+  
+  // Connection status is automatically synced to store
+  const status = useAppStore((s) => s.connection.status);
+  
+  return <div>Status: {status}</div>;
+}
+```
+
+## Testing
+
+To test reconnection behavior:
+
+1. Start the application
+2. Open browser DevTools → Network tab
+3. Observe WebSocket connection
+4. Stop the backend server
+5. Watch the client attempt to reconnect with increasing delays
+6. Restart the backend server
+7. Verify the client reconnects and re-subscribes to tasks
+
+## Best Practices
+
+1. **Always use the provided methods**: Use `subscribe()`/`unsubscribe()` instead of sending raw messages
+2. **Handle connection states in UI**: Show appropriate feedback for `connecting`, `reconnecting`, and `error` states
+3. **Clean up subscriptions**: Always unsubscribe when components unmount
+4. **Configure appropriately**: Adjust reconnect settings based on your use case
+5. **Monitor console logs**: Reconnection attempts are logged for debugging
+
@@ -2,7 +2,23 @@ import type { BackendMessageMap, BackendMessageType } from '@/lib/types/backend'

 type MessageHandler<T extends BackendMessageType> = (message: BackendMessageMap[T]) => void;

-type WebSocketStatus = 'idle' | 'connecting' | 'open' | 'closed' | 'error';
+type WebSocketStatus = 'idle' | 'connecting' | 'open' | 'closed' | 'error' | 'reconnecting';
+
+export interface ReconnectOptions {
+  enabled?: boolean;
+  maxAttempts?: number;
+  initialDelay?: number;
+  maxDelay?: number;
+  backoffMultiplier?: number;
+}
+
+const DEFAULT_RECONNECT_OPTIONS: Required<ReconnectOptions> = {
+  enabled: true,
+  maxAttempts: 10,
+  initialDelay: 1000,
+  maxDelay: 30000,
+  backoffMultiplier: 1.5,
+};

 export class WebSocketClient {
  private socket: WebSocket | null = null;
@@ -13,11 +29,19 @@ export class WebSocketClient {
    { resolve: (payload: unknown) => void; reject: (error: Error) => void; timeout: ReturnType<typeof setTimeout> }
  >();
  private pendingQueue: string[] = [];
+  private reconnectOptions: Required<ReconnectOptions>;
+  private reconnectAttempts = 0;
+  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+  private intentionalClose = false;
+  private subscriptions = new Set<string>();

  constructor(
    private url: string,
-    private onStatusChange?: (status: WebSocketStatus) => void
-  ) {}
+    private onStatusChange?: (status: WebSocketStatus) => void,
+    reconnectOptions?: ReconnectOptions
+  ) {
+    this.reconnectOptions = { ...DEFAULT_RECONNECT_OPTIONS, ...reconnectOptions };
+  }

  getStatus() {
    return this.status;
@@ -25,50 +49,63 @@ export class WebSocketClient {

  connect() {
    if (this.socket) return;
+    this.intentionalClose = false;
+    this.clearReconnectTimer();
    this.setStatus('connecting');
    this.socket = new WebSocket(this.url);

    this.socket.onopen = () => {
+      this.reconnectAttempts = 0;
      this.setStatus('open');
      this.flushQueue();
+      this.resubscribe();
    };

    this.socket.onmessage = (event) => {
-      try {
-        const message = JSON.parse(event.data) as BackendMessageMap[BackendMessageType];
-        if (message.type === 'response' && message.id) {
-          const pending = this.pendingRequests.get(message.id);
-          if (pending) {
-            clearTimeout(pending.timeout);
-            this.pendingRequests.delete(message.id);
-            pending.resolve(message.payload);
-          }
-          return;
-        }
-        if (message.type === 'error' && message.id) {
-          const pending = this.pendingRequests.get(message.id);
-          if (pending) {
-            clearTimeout(pending.timeout);
-            this.pendingRequests.delete(message.id);
-            const errorMessage =
-              typeof message.payload === 'object' && message.payload && 'message' in message.payload
-                ? String((message.payload as { message?: string }).message)
-                : 'WebSocket request failed';
-            pending.reject(new Error(errorMessage));
-          }
-          return;
-        }
+      // Backend batches multiple messages with newlines, so we need to split and handle each
+      const parts = (event.data as string).split('\n');
+      for (const part of parts) {
+        const trimmed = part.trim();
+        if (!trimmed) continue;

-        if (message.type !== 'notification') {
-          return;
+        try {
+          const message = JSON.parse(trimmed) as BackendMessageMap[BackendMessageType];
+          if (message.type === 'response' && (message as { id?: string }).id) {
+            const msgId = (message as { id: string }).id;
+            const pending = this.pendingRequests.get(msgId);
+            if (pending) {
+              clearTimeout(pending.timeout);
+              this.pendingRequests.delete(msgId);
+              pending.resolve(message.payload);
+            }
+            continue;
+          }
+          if (message.type === 'error' && (message as { id?: string }).id) {
+            const msgId = (message as { id: string }).id;
+            const pending = this.pendingRequests.get(msgId);
+            if (pending) {
+              clearTimeout(pending.timeout);
+              this.pendingRequests.delete(msgId);
+              const errorMessage =
+                typeof message.payload === 'object' && message.payload && 'message' in message.payload
+                  ? String((message.payload as { message?: string }).message)
+                  : 'WebSocket request failed';
+              pending.reject(new Error(errorMessage));
+            }
+            continue;
+          }
+
+          if (message.type !== 'notification') {
+            continue;
+          }
+          const action = (message as { action?: string })?.action as BackendMessageType | undefined;
+          if (!action) continue;
+          const handlers = this.handlers.get(action);
+          if (!handlers) continue;
+          handlers.forEach((handler) => handler(message));
+        } catch {
+          // Ignore parse errors for individual messages
        }
-        const action = message?.action;
-        if (!action) return;
-        const handlers = this.handlers.get(action);
-        if (!handlers) return;
-        handlers.forEach((handler) => handler(message));
-      } catch {
-        // Ignore malformed messages for now.
      }
    };

@@ -76,17 +113,21 @@ export class WebSocketClient {
      this.setStatus('error');
    };

-    this.socket.onclose = () => {
-      this.setStatus('closed');
+    this.socket.onclose = (event) => {
      this.socket = null;
+      this.handleDisconnect(event);
    };
  }

  disconnect() {
-    if (!this.socket) return;
-    this.socket.close();
-    this.socket = null;
+    this.intentionalClose = true;
+    this.clearReconnectTimer();
+    if (this.socket) {
+      this.socket.close();
+      this.socket = null;
+    }
    this.setStatus('closed');
+    this.cleanupPendingRequests();
  }

  send(payload: unknown) {
@@ -105,11 +146,35 @@ export class WebSocketClient {
        this.pendingRequests.delete(id);
        reject(new Error(`WebSocket request timed out: ${action}`));
      }, timeoutMs);
-      this.pendingRequests.set(id, { resolve, reject, timeout });
+      this.pendingRequests.set(id, { resolve: resolve as (payload: unknown) => void, reject, timeout });
      this.send({ id, type: 'request', action, payload });
    });
  }

+  subscribe(taskId: string) {
+    this.subscriptions.add(taskId);
+    if (this.status === 'open') {
+      this.send({
+        id: crypto.randomUUID(),
+        type: 'request',
+        action: 'task.subscribe',
+        payload: { task_id: taskId },
+      });
+    }
+  }
+
+  unsubscribe(taskId: string) {
+    this.subscriptions.delete(taskId);
+    if (this.status === 'open') {
+      this.send({
+        id: crypto.randomUUID(),
+        type: 'request',
+        action: 'task.unsubscribe',
+        payload: { task_id: taskId },
+      });
+    }
+  }
+
  on<T extends BackendMessageType>(type: T, handler: MessageHandler<T>) {
    const handlers = this.handlers.get(type) ?? new Set();
    handlers.add(handler as MessageHandler<BackendMessageType>);
@@ -126,6 +191,74 @@ export class WebSocketClient {
    }
  }

+  private handleDisconnect(event: CloseEvent) {
+    this.setStatus('closed');
+
+    // Don't reconnect if this was an intentional close
+    if (this.intentionalClose) {
+      return;
+    }
+
+    // Don't reconnect if reconnect is disabled
+    if (!this.reconnectOptions.enabled) {
+      return;
+    }
+
+    // Don't reconnect if we've exceeded max attempts
+    if (this.reconnectAttempts >= this.reconnectOptions.maxAttempts) {
+      console.warn(`WebSocket max reconnect attempts (${this.reconnectOptions.maxAttempts}) reached`);
+      this.setStatus('error');
+      this.cleanupPendingRequests();
+      return;
+    }
+
+    // Calculate delay with exponential backoff
+    const delay = Math.min(
+      this.reconnectOptions.initialDelay * Math.pow(this.reconnectOptions.backoffMultiplier, this.reconnectAttempts),
+      this.reconnectOptions.maxDelay
+    );
+
+    this.reconnectAttempts++;
+    this.setStatus('reconnecting');
+
+    console.log(
+      `WebSocket disconnected (code: ${event.code}, reason: ${event.reason || 'none'}). ` +
+        `Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts}/${this.reconnectOptions.maxAttempts})...`
+    );
+
+    this.reconnectTimer = setTimeout(() => {
+      this.connect();
+    }, delay);
+  }
+
+  private clearReconnectTimer() {
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
+  }
+
+  private cleanupPendingRequests() {
+    // Reject all pending requests
+    this.pendingRequests.forEach(({ reject, timeout }) => {
+      clearTimeout(timeout);
+      reject(new Error('WebSocket connection closed'));
+    });
+    this.pendingRequests.clear();
+  }
+
+  private resubscribe() {
+    // Re-subscribe to all tasks after reconnection
+    this.subscriptions.forEach((taskId) => {
+      this.send({
+        id: crypto.randomUUID(),
+        type: 'request',
+        action: 'task.subscribe',
+        payload: { task_id: taskId },
+      });
+    });
+  }
+
  private flushQueue() {
    if (!this.socket || this.status !== 'open') return;
    this.pendingQueue.forEach((data) => this.socket?.send(data));
@@ -310,5 +310,20 @@ export function registerWsHandlers(store: StoreApi<AppState>) {
    'system.error': () => {
      // TODO: surface as toast/notification once UI is ready.
    },
+    'comment.added': (message: BackendMessageMap['comment.added']) => {
+      const payload = message.payload;
+      console.log('[WS] comment.added payload:', JSON.stringify(payload, null, 2));
+      store.getState().addComment({
+        id: payload.comment_id,
+        task_id: payload.task_id,
+        author_type: payload.author_type,
+        author_id: payload.author_id,
+        content: payload.content,
+        type: (payload.type as 'message' | 'content' | 'tool_call' | 'progress' | 'error' | 'status') || 'message',
+        metadata: payload.metadata,
+        requests_input: payload.requests_input,
+        created_at: payload.created_at,
+      });
+    },
  };
 }
@@ -11,25 +11,38 @@ export function useWebSocket(store: StoreApi<AppState>, url: string) {
  const clientRef = useRef<WebSocketClient | null>(null);

  useEffect(() => {
-    const client = new WebSocketClient(url, (status) => {
-      const setConnectionStatus = store.getState().setConnectionStatus;
-      switch (status) {
-        case 'connecting':
-          setConnectionStatus('connecting', null);
-          break;
-        case 'open':
-          setConnectionStatus('connected', null);
-          break;
-        case 'error':
-          setConnectionStatus('error', 'WebSocket error');
-          break;
-        case 'closed':
-        case 'idle':
-        default:
-          setConnectionStatus('disconnected', null);
-          break;
+    const client = new WebSocketClient(
+      url,
+      (status) => {
+        const setConnectionStatus = store.getState().setConnectionStatus;
+        switch (status) {
+          case 'connecting':
+            setConnectionStatus('connecting', null);
+            break;
+          case 'open':
+            setConnectionStatus('connected', null);
+            break;
+          case 'reconnecting':
+            setConnectionStatus('reconnecting', null);
+            break;
+          case 'error':
+            setConnectionStatus('error', 'WebSocket connection failed');
+            break;
+          case 'closed':
+          case 'idle':
+          default:
+            setConnectionStatus('disconnected', null);
+            break;
+        }
+      },
+      {
+        enabled: true,
+        maxAttempts: 10,
+        initialDelay: 1000,
+        maxDelay: 30000,
+        backoffMultiplier: 1.5,
      }
-    });
+    );
    clientRef.current = client;
    client.connect();
    setWebSocketClient(client);
@@ -24,7 +24,10 @@
    "radix-ui": "^1.4.3",
    "react": "19.2.3",
    "react-dom": "19.2.3",
+    "react-markdown": "^10.1.0",
    "react-resizable-panels": "^2.1.9",
+    "remark-breaks": "^4.0.0",
+    "remark-gfm": "^4.0.1",
    "shadcn": "^3.6.3",
    "tailwind-merge": "^3.4.0",
    "tw-animate-css": "^1.4.0",