@assistkick/create 1.9.0 → 1.11.0

package/templates/assistkick-product-system/packages/backend/src/services/chat_ws_handler.ts

@@ -0,0 +1,894 @@
+/**
+ * WebSocket handler for Chat v2 message streaming.
+ * Authenticates via cookie, receives user messages, spawns Claude CLI
+ * via ChatCliBridge, and forwards parsed stream-json events in real time.
+ * Persists messages to DB on successful stream completion; discards on cancel.
+ *
+ * Client → Server messages:
+ *   { type: 'send_message', message, projectId, sessionId, claudeSessionId,
+ *     isNewSession, permissionMode?, allowedTools?, permissionPromptTool? }
+ *   { type: 'permission_response', requestId, decision }
+ *   { type: 'stop_stream' }
+ *
+ * Server → Client messages:
+ *   { type: 'stream_start', claudeSessionId }
+ *   { type: 'stream_event', event }    — forwarded parsed stream-json event
+ *   { type: 'stream_end', exitCode, claudeSessionId }
+ *   { type: 'stream_cancelled', claudeSessionId }  — user-initiated stop
+ *   { type: 'stream_error', error }
+ *   { type: 'permission_request', requestId, toolName, input }
+ *   { type: 'error', message }
+ */
+
+import type { IncomingMessage } from 'node:http';
+import type { Duplex } from 'node:stream';
+import type { WebSocketServer, WebSocket } from 'ws';
+import type { AuthService } from './auth_service.js';
+import type { ChatCliBridge, PermissionMode } from './chat_cli_bridge.js';
+import type { PermissionService, PermissionDecision, PermissionRequest } from './permission_service.js';
+import type { ChatMessageRepository } from './chat_message_repository.js';
+import type { ChatSessionService } from './chat_session_service.js';
+import type { TitleGeneratorService } from './title_generator_service.js';
+
+export interface ChatWsHandlerDeps {
+  wss: WebSocketServer;
+  authService: AuthService;
+  chatCliBridge: ChatCliBridge;
+  permissionService?: PermissionService;
+  chatMessageRepository?: ChatMessageRepository;
+  chatSessionService?: ChatSessionService;
+  titleGeneratorService?: TitleGeneratorService;
+  log: (tag: string, ...args: unknown[]) => void;
+}
+
+export interface SendMessagePayload {
+  type: 'send_message';
+  message: string;
+  projectId: string;
+  /** The chat session ID (csess_xxx) for message persistence. */
+  sessionId?: string;
+  claudeSessionId: string;
+  isNewSession: boolean;
+  permissionMode?: PermissionMode;
+  allowedTools?: string[];
+  /** Relative workspace paths of attached files to include in the prompt */
+  attachments?: string[];
+  /** Attachment content blocks for DB persistence alongside the user message */
+  attachmentBlocks?: Array<{ type: 'image' | 'file'; path: string; name: string; mimeType: string }>;
+}
+
+export interface PermissionResponsePayload {
+  type: 'permission_response';
+  requestId: string;
+  projectId: string;
+  decision: PermissionDecision;
+}
+
+export interface StopStreamPayload {
+  type: 'stop_stream';
+  /** Target a specific Claude session. If omitted, stops all active streams on this WS. */
+  claudeSessionId?: string;
+}
+
+export interface AttachSessionPayload {
+  type: 'attach_session';
+  claudeSessionId: string;
+}
+
+export type ChatClientMessage = SendMessagePayload | PermissionResponsePayload | StopStreamPayload | AttachSessionPayload;
+
+export interface StreamStartMessage {
+  type: 'stream_start';
+  claudeSessionId: string;
+}
+
+export interface StreamEventMessage {
+  type: 'stream_event';
+  event: Record<string, unknown>;
+}
+
+export interface StreamEndMessage {
+  type: 'stream_end';
+  exitCode: number | null;
+  claudeSessionId: string;
+  /** Error details when exitCode is non-zero (stderr + non-JSON stdout from CLI) */
+  error?: string;
+}
+
+export interface StreamCancelledMessage {
+  type: 'stream_cancelled';
+  claudeSessionId: string;
+}
+
+export interface StreamErrorMessage {
+  type: 'stream_error';
+  error: string;
+}
+
+export interface PermissionRequestMessage {
+  type: 'permission_request';
+  requestId: string;
+  toolName: string;
+  input: Record<string, unknown>;
+}
+
+export interface StreamResumedMessage {
+  type: 'stream_resumed';
+  claudeSessionId: string;
+}
+
+export interface SystemPromptMessage {
+  type: 'system_prompt';
+  prompt: string;
+}
+
+export interface SessionRenamedMessage {
+  type: 'session_renamed';
+  sessionId: string;
+  name: string;
+}
+
+export interface ErrorMessage {
+  type: 'error';
+  message: string;
+}
+
+export type ChatServerMessage =
+  | StreamStartMessage
+  | StreamEventMessage
+  | StreamEndMessage
+  | StreamCancelledMessage
+  | StreamErrorMessage
+  | StreamResumedMessage
+  | PermissionRequestMessage
+  | SystemPromptMessage
+  | SessionRenamedMessage
+  | ErrorMessage;
+
+/** Tracks in-flight stream context for message persistence. */
+interface StreamContext {
+  /** The chat session ID (csess_xxx) — null if frontend didn't provide it. */
+  sessionId: string | null;
+  /** The original user message text. */
+  userMessage: string;
+  /** Attachment content blocks from the frontend (image/file metadata). */
+  attachmentBlocks: Array<{ type: 'image' | 'file'; path: string; name: string; mimeType: string }>;
+  /** The last known assistant content blocks (JSON string), accumulated from stream events. */
+  lastAssistantContent: string;
+  /** Whether messages have been persisted (on WS disconnect). */
+  persisted: boolean;
+  /** ID of the persisted assistant message, set after partial persistence completes. */
+  persistedAssistantMessageId: string | null;
+  /** Promise that resolves when partial persistence is complete (for sequencing with final persist). */
+  persistPromise: Promise<void> | null;
+  /** Last known context usage JSON from the most recent assistant event (per-call, not cumulative). */
+  lastContextUsageJson: string | null;
+  /** Context window size from the model's modelUsage (extracted from result event). */
+  contextWindow: number | null;
+}
+
+const parseCookies = (cookieHeader: string | undefined): Record<string, string> => {
+  if (!cookieHeader) return {};
+  const cookies: Record<string, string> = {};
+  for (const pair of cookieHeader.split(';')) {
+    const [key, ...rest] = pair.trim().split('=');
+    if (key) {
+      cookies[key] = rest.join('=');
+    }
+  }
+  return cookies;
+};
+
+/**
+ * Extract the content array from a stream-json event for persistence.
+ * Returns the parsed content blocks, or null if not extractable.
+ */
+const extractEventContent = (event: Record<string, unknown>): Array<Record<string, unknown>> | null => {
+  const message = event.message as Record<string, unknown> | undefined;
+  if (!message) return null;
+  const content = message.content;
+  if (!Array.isArray(content)) return null;
+  return content as Array<Record<string, unknown>>;
+};
+
+/**
+ * Merge new assistant content blocks with accumulated content, preserving
+ * tool_use and tool_result blocks from previous turns.
+ * Within a turn, assistant events are cumulative (replace current turn blocks).
+ * Across turns (detected by missing tool_use IDs), previous content is preserved.
+ */
+const mergeAssistantContent = (
+  existing: Array<Record<string, unknown>>,
+  newBlocks: Array<Record<string, unknown>>,
+): Array<Record<string, unknown>> => {
+  const existingToolUseIds = new Set(
+    existing.filter(b => b.type === 'tool_use' && typeof b.id === 'string').map(b => b.id as string),
+  );
+  const newToolUseIds = new Set(
+    newBlocks.filter(b => b.type === 'tool_use' && typeof b.id === 'string').map(b => b.id as string),
+  );
+
+  const isNewTurn = existingToolUseIds.size > 0 &&
+    ![...existingToolUseIds].some(id => newToolUseIds.has(id));
+
+  if (isNewTurn) {
+    return [...existing, ...newBlocks];
+  }
+  // Same turn: replace with cumulative snapshot, preserving tool_result blocks
+  const existingToolResults = existing.filter(b => b.type === 'tool_result');
+  return [...newBlocks, ...existingToolResults];
+};
+
+/**
+ * Extract tool_result blocks from a human/tool_result event.
+ */
+const extractToolResults = (event: Record<string, unknown>): Array<Record<string, unknown>> => {
+  if (event.type !== 'tool_result' && event.type !== 'human' && event.type !== 'user') return [];
+  const content = extractEventContent(event);
+  if (!content) return [];
+  return content.filter(b => b.type === 'tool_result');
+};
+
+export class ChatWsHandler {
+  private readonly wss: WebSocketServer;
+  private readonly authService: AuthService;
+  private readonly chatCliBridge: ChatCliBridge;
+  private readonly permissionService: PermissionService | null;
+  private readonly chatMessageRepository: ChatMessageRepository | null;
+  private readonly chatSessionService: ChatSessionService | null;
+  private readonly titleGeneratorService: TitleGeneratorService | null;
+  private readonly log: ChatWsHandlerDeps['log'];
+  private readonly activeStreams = new Map<WebSocket, Set<string>>();
+
+  /** Map Claude session IDs to their active WebSocket for permission routing */
+  private readonly sessionToWs = new Map<string, WebSocket>();
+
+  /** Track session IDs that were cancelled by user (stop_stream) so completion sends stream_cancelled instead of stream_end */
+  private readonly cancelledStreams = new Set<string>();
+
+  /** Track stream context per Claude session ID for message persistence. */
+  private readonly streamContexts = new Map<string, StreamContext>();
+
+  constructor({ wss, authService, chatCliBridge, permissionService, chatMessageRepository, chatSessionService, titleGeneratorService, log }: ChatWsHandlerDeps) {
+    this.wss = wss;
+    this.authService = authService;
+    this.chatCliBridge = chatCliBridge;
+    this.permissionService = permissionService || null;
+    this.chatMessageRepository = chatMessageRepository || null;
+    this.chatSessionService = chatSessionService || null;
+    this.titleGeneratorService = titleGeneratorService || null;
+    this.log = log;
+
+    // Register the permission request handler to route requests to the right WebSocket
+    if (this.permissionService) {
+      this.permissionService.setPermissionRequestHandler(this.onPermissionRequest);
+    }
+  }
+
+  /** Remove a single Claude session from a WebSocket's active stream set. */
+  private removeActiveStream = (ws: WebSocket, claudeSessionId: string): void => {
+    const streams = this.activeStreams.get(ws);
+    if (streams) {
+      streams.delete(claudeSessionId);
+      if (streams.size === 0) this.activeStreams.delete(ws);
+    }
+  };
+
+  handleUpgrade = (req: IncomingMessage, socket: Duplex, head: Buffer): void => {
+    const url = new URL(req.url || '', 'http://localhost');
+    if (url.pathname !== '/api/chat') {
+      return;
+    }
+
+    this.wss.handleUpgrade(req, socket, head, (ws) => {
+      this.onConnection(ws, req);
+    });
+  };
+
+  private sendWsMessage = (ws: WebSocket, msg: ChatServerMessage): void => {
+    if (ws.readyState === ws.OPEN) {
+      ws.send(JSON.stringify(msg));
+    }
+  };
+
+  private onConnection = async (ws: WebSocket, req: IncomingMessage): Promise<void> => {
+    const cookies = parseCookies(req.headers.cookie);
+    const token = cookies['access_token'];
+
+    if (!token) {
+      this.log('CHAT_WS', 'Connection rejected — no access token');
+      ws.close(4001, 'Unauthorized');
+      return;
+    }
+
+    let payload: { sub: string; email?: string; role?: string };
+    try {
+      payload = await this.authService.verifyToken(token);
+    } catch {
+      this.log('CHAT_WS', 'Connection rejected — invalid token');
+      ws.close(4001, 'Invalid token');
+      return;
+    }
+
+    if (payload.role !== 'admin') {
+      this.log('CHAT_WS', `Connection rejected — user ${payload.sub} is not admin`);
+      ws.close(4003, 'Forbidden');
+      return;
+    }
+
+    this.log('CHAT_WS', `Admin ${payload.email} connected`);
+
+    ws.on('message', (raw: Buffer | string) => {
+      let msg: ChatClientMessage;
+      try {
+        msg = JSON.parse(typeof raw === 'string' ? raw : raw.toString());
+      } catch {
+        this.sendWsMessage(ws, { type: 'error', message: 'Invalid JSON' });
+        return;
+      }
+
+      if (msg.type === 'send_message') {
+        this.handleSendMessage(ws, msg);
+      } else if (msg.type === 'permission_response') {
+        this.handlePermissionResponse(msg);
+      } else if (msg.type === 'stop_stream') {
+        this.handleStopStream(ws, msg);
+      } else if (msg.type === 'attach_session') {
+        this.handleAttachSession(ws, msg);
+      }
+    });
+
+    ws.on('close', () => {
+      this.log('CHAT_WS', `Admin ${payload.email} disconnected`);
+      this.cancelActiveStream(ws);
+    });
+
+    ws.on('error', (err) => {
+      this.log('CHAT_WS', `WebSocket error: ${err.message}`);
+      this.cancelActiveStream(ws);
+    });
+  };
+
+  /**
+   * Format the user message with attachment file path references appended.
+   */
+  private formatMessageWithAttachments = (message: string, attachments?: string[]): string => {
+    if (!attachments || attachments.length === 0) return message;
+    const fileList = attachments.map((p) => `- ${p}`).join('\n');
+    return `${message}\n\nAttached files:\n${fileList}`;
+  };
+
+  private handleSendMessage = async (ws: WebSocket, msg: SendMessagePayload): Promise<void> => {
+    const { message, projectId, claudeSessionId, isNewSession } = msg;
+
+    // Reject if this specific Claude session is already streaming
+    const wsStreams = this.activeStreams.get(ws);
+    if (wsStreams?.has(claudeSessionId)) {
+      this.sendWsMessage(ws, {
+        type: 'error',
+        message: 'A message is already being processed. Wait for it to complete.',
+      });
+      return;
+    }
+
+    // Validate required fields
+    if (!message || !projectId || !claudeSessionId) {
+      this.sendWsMessage(ws, {
+        type: 'error',
+        message: 'Missing required fields: message, projectId, claudeSessionId',
+      });
+      return;
+    }
+
+    let activeSet = this.activeStreams.get(ws);
+    if (!activeSet) {
+      activeSet = new Set();
+      this.activeStreams.set(ws, activeSet);
+    }
+    activeSet.add(claudeSessionId);
+    this.sessionToWs.set(claudeSessionId, ws);
+
+    // Track stream context for persistence (save on completion, discard on cancel)
+    this.streamContexts.set(claudeSessionId, {
+      sessionId: msg.sessionId || null,
+      userMessage: message,
+      attachmentBlocks: msg.attachmentBlocks || [],
+      lastAssistantContent: '[]',
+      persisted: false,
+      persistedAssistantMessageId: null,
+      persistPromise: null,
+      lastContextUsageJson: null,
+      contextWindow: null,
+    });
+
+    this.sendWsMessage(ws, { type: 'stream_start', claudeSessionId });
+    this.log('CHAT_WS', `Streaming started for session ${claudeSessionId}`);
+
+    const permissionMode = msg.permissionMode || 'skip';
+    const formattedMessage = this.formatMessageWithAttachments(message, msg.attachments);
+
+    // Look up continuation context from session DB when starting a new session
+    let continuationContext: string | undefined;
+    if (isNewSession && msg.sessionId && this.chatSessionService) {
+      const session = await this.chatSessionService.getSession(msg.sessionId);
+      if (session?.continuationContext) {
+        continuationContext = session.continuationContext;
+      }
+    }
+
+    try {
+      const { completion, systemPrompt } = this.chatCliBridge.spawn({
+        projectId,
+        message: formattedMessage,
+        claudeSessionId,
+        isNewSession,
+        permissionMode,
+        allowedTools: msg.allowedTools,
+        continuationContext,
+        onEvent: (event) => {
+          // Forward to current WS (may have been re-attached after disconnect)
+          const currentWs = this.sessionToWs.get(claudeSessionId);
+          if (currentWs) {
+            this.sendWsMessage(currentWs, { type: 'stream_event', event });
+          }
+
+          // Accumulate assistant content and context usage for persistence
+          const ctx = this.streamContexts.get(claudeSessionId);
+          if (!ctx) return;
+
+          // Extract per-call context usage from assistant events only.
+          // The result event contains CUMULATIVE usage across all turns — not the context window size.
+          if (event.type === 'assistant') {
+            const message = event.message as Record<string, unknown> | undefined;
+            const usage = message?.usage as Record<string, unknown> | undefined;
+            if (usage) {
+              const input = typeof usage.input_tokens === 'number' ? usage.input_tokens : 0;
+              const output = typeof usage.output_tokens === 'number' ? usage.output_tokens : 0;
+              const cacheCreation = typeof usage.cache_creation_input_tokens === 'number' ? usage.cache_creation_input_tokens : 0;
+              const cacheRead = typeof usage.cache_read_input_tokens === 'number' ? usage.cache_read_input_tokens : 0;
+              ctx.lastContextUsageJson = JSON.stringify({
+                inputTokens: input,
+                outputTokens: output,
+                cacheCreationTokens: cacheCreation,
+                cacheReadTokens: cacheRead,
+                // totalTokens = input context only (output tokens don't count toward context window)
+                totalTokens: input + cacheCreation + cacheRead,
+                contextWindow: ctx.contextWindow,
+              });
+            }
+          }
+
+          // Extract contextWindow from result event's modelUsage
+          if (event.type === 'result') {
+            const modelUsage = event.modelUsage as Record<string, Record<string, unknown>> | undefined;
+            if (modelUsage) {
+              for (const info of Object.values(modelUsage)) {
+                if (typeof info.contextWindow === 'number') {
+                  ctx.contextWindow = info.contextWindow;
+                  // Re-emit the last context usage JSON with the contextWindow included
+                  if (ctx.lastContextUsageJson) {
+                    try {
+                      const parsed = JSON.parse(ctx.lastContextUsageJson);
+                      parsed.contextWindow = ctx.contextWindow;
+                      ctx.lastContextUsageJson = JSON.stringify(parsed);
+                    } catch { /* ignore */ }
+                  }
+                  break;
+                }
+              }
+            }
+          }
+
+          if (event.type === 'assistant') {
+            const newBlocks = extractEventContent(event);
+            if (newBlocks && newBlocks.length > 0) {
+              const existing = JSON.parse(ctx.lastAssistantContent) as Array<Record<string, unknown>>;
+              const merged = mergeAssistantContent(existing, newBlocks);
+              ctx.lastAssistantContent = JSON.stringify(merged);
+            }
+          } else {
+            // Capture tool_result blocks from human/tool_result events
+            const toolResults = extractToolResults(event);
+            if (toolResults.length > 0) {
+              const existing = JSON.parse(ctx.lastAssistantContent) as Array<Record<string, unknown>>;
+              const existingToolResultIds = new Set(
+                existing.filter(b => b.type === 'tool_result').map(b => b.tool_use_id as string),
+              );
+              const newResults = toolResults.filter(b => !existingToolResultIds.has(b.tool_use_id as string));
+              if (newResults.length > 0) {
+                ctx.lastAssistantContent = JSON.stringify([...existing, ...newResults]);
+              }
+            }
+          }
+        },
+      });
+
+      // Send composed system prompt to frontend and persist to session
+      if (systemPrompt) {
+        this.sendWsMessage(ws, { type: 'system_prompt', prompt: systemPrompt });
+
+        if (msg.sessionId && this.chatSessionService) {
+          this.chatSessionService.updateSystemPrompt(msg.sessionId, systemPrompt)
+            .catch((err: Error) => {
+              this.log('CHAT_WS', `Failed to persist system prompt for session ${msg.sessionId}: ${err.message}`);
+            });
+        }
+      }
+
+      completion
+        .then((result) => {
+          // Look up the current WS — may differ from the original if client reconnected
+          const currentWs = this.sessionToWs.get(claudeSessionId);
+          if (currentWs) {
+            this.removeActiveStream(currentWs, claudeSessionId);
+          }
+          this.sessionToWs.delete(claudeSessionId);
+
+          // If this was a user-initiated stop, send stream_cancelled instead of stream_end
+          if (this.cancelledStreams.has(claudeSessionId)) {
+            this.cancelledStreams.delete(claudeSessionId);
+            this.streamContexts.delete(claudeSessionId);
+            if (currentWs) {
+              this.sendWsMessage(currentWs, {
+                type: 'stream_cancelled',
+                claudeSessionId,
+              });
+            }
+            this.log('CHAT_WS', `Stream cancelled by user for session ${claudeSessionId}`);
+          } else {
+            // Capture stream context for title generation before persistMessages deletes it
+            const titleCtx = this.streamContexts.get(claudeSessionId);
+            const titleSessionId = titleCtx?.sessionId ?? null;
+            const titleUserMessage = titleCtx?.userMessage ?? '';
+            const titleAssistantContent = titleCtx?.lastAssistantContent ?? '[]';
+
+            // Persist messages on successful completion (handles both fresh and partial-persist cases)
+            this.persistMessages(claudeSessionId);
+
+            // Fire-and-forget auto-title generation
+            if (titleSessionId) {
+              this.maybeGenerateTitle(claudeSessionId, titleSessionId, titleUserMessage, titleAssistantContent);
+            }
+
+            // Build error details for non-zero exit codes
+            let errorDetail: string | undefined;
+            if (result.exitCode !== 0 && result.exitCode !== null) {
+              const parts: string[] = [];
+              if (result.stderr) parts.push(result.stderr);
+              if (result.stdoutNonJsonLines) parts.push(result.stdoutNonJsonLines);
+              errorDetail = parts.length > 0
+                ? parts.join('\n')
+                : `CLI exited with code ${result.exitCode}`;
+            }
+
+            if (currentWs) {
+              const msg: StreamEndMessage = {
+                type: 'stream_end',
+                exitCode: result.exitCode,
+                claudeSessionId: result.claudeSessionId,
+              };
+              if (errorDetail) msg.error = errorDetail;
+              this.sendWsMessage(currentWs, msg);
+            }
+            this.log('CHAT_WS', `Streaming ended for session ${claudeSessionId}, exit code ${result.exitCode}${errorDetail ? `, error: ${errorDetail}` : ''}`);
+          }
+        })
+        .catch((err: Error) => {
+          const currentWs = this.sessionToWs.get(claudeSessionId);
+          if (currentWs) {
+            this.removeActiveStream(currentWs, claudeSessionId);
+          }
+          this.sessionToWs.delete(claudeSessionId);
+
+          // Persist whatever we have even on error
+          this.persistMessages(claudeSessionId);
+
+          if (currentWs) {
+            this.sendWsMessage(currentWs, {
+              type: 'stream_error',
+              error: err.message,
+            });
+          }
+          this.log('CHAT_WS', `Streaming error for session ${claudeSessionId}: ${err.message}`);
+        });
+    } catch (err: unknown) {
+      this.removeActiveStream(ws, claudeSessionId);
+      this.sessionToWs.delete(claudeSessionId);
+      this.streamContexts.delete(claudeSessionId);
+      const errorMsg = err instanceof Error ? err.message : 'Unknown error';
+      this.sendWsMessage(ws, { type: 'stream_error', error: errorMsg });
+      this.log('CHAT_WS', `Failed to spawn CLI for session ${claudeSessionId}: ${errorMsg}`);
+    }
+  };
+
+  /**
+   * Persist the user message and assistant response to the database.
+   * Handles two cases:
+   * - Fresh: no prior persistence — save both messages.
+   * - Partial: messages were already saved on WS disconnect — update the assistant message with final content.
+   */
+  private persistMessages = (claudeSessionId: string): void => {
+    const ctx = this.streamContexts.get(claudeSessionId);
+    this.streamContexts.delete(claudeSessionId);
+
+    if (!ctx || !ctx.sessionId || !this.chatMessageRepository) return;
+
+    // Persist context usage to the session
+    if (ctx.lastContextUsageJson && this.chatSessionService) {
+      this.chatSessionService.updateContextUsage(ctx.sessionId, ctx.lastContextUsageJson)
+        .catch((err: Error) => {
+          this.log('CHAT_WS', `Failed to persist context usage for session ${ctx.sessionId}: ${err.message}`);
+        });
+    }
+
+    const { sessionId, lastAssistantContent } = ctx;
+
+    if (ctx.persisted) {
+      // Messages were already saved on WS disconnect — update the assistant message with final content
+      const doUpdate = () => {
+        if (ctx.persistedAssistantMessageId) {
+          this.chatMessageRepository!.updateMessageContent(ctx.persistedAssistantMessageId, lastAssistantContent)
+            .then(() => {
+              this.log('CHAT_WS', `Updated assistant message ${ctx.persistedAssistantMessageId} with final content for session ${sessionId}`);
+            })
+            .catch((err: Error) => {
+              this.log('CHAT_WS', `Failed to update assistant message for session ${sessionId}: ${err.message}`);
+            });
+        }
+      };
+
+      // Wait for partial persist to complete before updating
+      if (ctx.persistPromise) {
+        ctx.persistPromise.then(doUpdate);
+      } else {
+        doUpdate();
+      }
+    } else {
+      // Fresh path — save both messages
+      const userBlocks: Array<Record<string, unknown>> = [{ type: 'text', text: ctx.userMessage }];
+      for (const block of ctx.attachmentBlocks) {
+        userBlocks.push({ type: block.type, path: block.path, name: block.name, mimeType: block.mimeType });
+      }
+      const userContent = JSON.stringify(userBlocks);
+
+      this.chatMessageRepository.getNextOrderIndex(sessionId)
+        .then((nextIndex) => {
+          return Promise.all([
+            this.chatMessageRepository!.saveMessage(sessionId, 'user', userContent, nextIndex),
+            this.chatMessageRepository!.saveMessage(sessionId, 'assistant', lastAssistantContent, nextIndex + 1),
+          ]);
+        })
+        .catch((err: Error) => {
+          this.log('CHAT_WS', `Failed to persist messages for session ${sessionId}: ${err.message}`);
+        });
+    }
+  };
+
+  /**
+   * Persist partial messages immediately when a WebSocket disconnects mid-stream.
+   * Saves the user message and whatever assistant content has been accumulated so far.
+   */
+  private persistPartialMessages = (claudeSessionId: string): void => {
+    const ctx = this.streamContexts.get(claudeSessionId);
+    if (!ctx || !ctx.sessionId || !this.chatMessageRepository || ctx.persisted) return;
+
+    ctx.persisted = true;
+
+    const { sessionId, userMessage, lastAssistantContent, attachmentBlocks } = ctx;
+    const userBlocks: Array<Record<string, unknown>> = [{ type: 'text', text: userMessage }];
+    for (const block of attachmentBlocks) {
+      userBlocks.push({ type: block.type, path: block.path, name: block.name, mimeType: block.mimeType });
+    }
+    const userContent = JSON.stringify(userBlocks);
+
+    ctx.persistPromise = this.chatMessageRepository.getNextOrderIndex(sessionId)
+      .then((nextIndex) => {
+        return Promise.all([
+          this.chatMessageRepository!.saveMessage(sessionId, 'user', userContent, nextIndex),
+          this.chatMessageRepository!.saveMessage(sessionId, 'assistant', lastAssistantContent, nextIndex + 1),
+        ]);
+      })
+      .then(([, assistantRow]) => {
+        ctx.persistedAssistantMessageId = assistantRow.id;
+        this.log('CHAT_WS', `Persisted partial messages for session ${sessionId} (assistant msg: ${assistantRow.id})`);
+      })
+      .catch((err: Error) => {
+        this.log('CHAT_WS', `Failed to persist partial messages for session ${sessionId}: ${err.message}`);
+      });
+  };
+
+  /**
+   * Fire-and-forget: attempt to auto-generate a title for a session that still has
+   * the default timestamp-based name. Runs after messages are persisted so the title
+   * can consider both the user message and assistant reply.
+   */
+  private maybeGenerateTitle = (claudeSessionId: string, sessionId: string, userMessage: string, assistantContent: string): void => {
+    if (!this.titleGeneratorService || !this.chatSessionService) return;
+
+    const titleGen = this.titleGeneratorService;
+    const sessionService = this.chatSessionService;
+
+    // Fire-and-forget — never blocks the main stream
+    sessionService.getSession(sessionId)
+      .then((session) => {
+        if (!session) return;
+        if (!titleGen.isDefaultName(session.name)) return;
+
+        return titleGen.generateTitle(userMessage, assistantContent)
+          .then((title) => {
+            if (!title) return;
+            return sessionService.renameSession(sessionId, title)
+              .then(() => {
+                // Push session_renamed to the connected WebSocket client
+                const ws = this.sessionToWs.get(claudeSessionId);
+                if (ws) {
+                  this.sendWsMessage(ws, {
+                    type: 'session_renamed',
+                    sessionId,
+                    name: title,
+                  });
+                }
+                // Also broadcast to all connected clients in case the sidebar is open on another tab
+                this.broadcastSessionRenamed(sessionId, title, claudeSessionId);
+              });
+          });
+      })
+      .catch((err: Error) => {
+        this.log('CHAT_WS', `Auto-title generation failed for session ${sessionId}: ${err.message}`);
+      });
+  };
+
+  /**
+   * Broadcast session_renamed to all connected WebSocket clients except the one
+   * already notified (identified by claudeSessionId).
+   */
+  private broadcastSessionRenamed = (sessionId: string, name: string, excludeClaudeSessionId: string): void => {
+    const excludeWs = this.sessionToWs.get(excludeClaudeSessionId);
+    for (const client of this.wss.clients) {
+      if (client !== excludeWs && client.readyState === client.OPEN) {
+        this.sendWsMessage(client as WebSocket, {
+          type: 'session_renamed',
+          sessionId,
+          name,
+        });
+      }
+    }
+  };
+
+  /**
+   * Handle permission_response from the frontend.
+   * Routes it to the PermissionService to resolve the pending request.
+   */
+  private handlePermissionResponse = (msg: PermissionResponsePayload): void => {
+    if (!this.permissionService) {
+      this.log('CHAT_WS', 'Permission response received but no permission service configured');
+      return;
+    }
+
+    const { requestId, projectId, decision } = msg;
+    if (!requestId || !decision) {
+      this.log('CHAT_WS', 'Permission response missing requestId or decision');
+      return;
+    }
+
+    this.permissionService.resolvePermission(requestId, projectId, decision).catch((err: Error) => {
+      this.log('CHAT_WS', `Failed to resolve permission: ${err.message}`);
+    });
+  };
+
+  /**
+   * Called by PermissionService when a new permission request arrives
+   * from the MCP server and needs to be forwarded to the frontend.
+   */
+  private onPermissionRequest = (request: PermissionRequest): void => {
+    const ws = this.sessionToWs.get(request.claudeSessionId);
+    if (!ws) {
+      this.log('CHAT_WS', `No WebSocket found for session ${request.claudeSessionId} — cannot send permission request`);
+      return;
+    }
+
+    this.sendWsMessage(ws, {
+      type: 'permission_request',
+      requestId: request.requestId,
+      toolName: request.toolName,
+      input: request.input,
+    });
+
+    this.log('CHAT_WS', `Forwarded permission request ${request.requestId} for tool "${request.toolName}"`);
+  };
+
+  /**
+   * Handle user-initiated stop: kill the CLI process but keep the WebSocket open.
+   * The completion handler will detect the cancellation and send stream_cancelled.
+   */
+  private handleStopStream = (ws: WebSocket, msg: StopStreamPayload): void => {
+    const streams = this.activeStreams.get(ws);
+    if (!streams || streams.size === 0) {
+      this.sendWsMessage(ws, { type: 'error', message: 'No active stream to stop' });
+      return;
+    }
+
+    // If a specific session was targeted, stop just that one; otherwise stop all
+    const targets = msg.claudeSessionId
+      ? (streams.has(msg.claudeSessionId) ? [msg.claudeSessionId] : [])
+      : [...streams];
+
+    if (targets.length === 0) {
+      this.sendWsMessage(ws, { type: 'error', message: 'No active stream to stop' });
+      return;
+    }
+
+    for (const sessionId of targets) {
+      this.cancelledStreams.add(sessionId);
+      this.chatCliBridge.kill(sessionId);
+
+      if (this.permissionService) {
+        this.permissionService.cancelPendingRequests(sessionId);
+      }
+
+      this.log('CHAT_WS', `User requested stop for session ${sessionId}`);
+    }
+  };
+
+  /**
+   * Handle WS disconnect: clean up WS mappings but let the CLI process continue running.
+   * Persists partial messages so they survive the disconnect.
+   * The completion handler will still run and finalize persistence.
+   */
+  private cancelActiveStream = (ws: WebSocket): void => {
+    const streams = this.activeStreams.get(ws);
+    if (streams && streams.size > 0) {
+      // Do NOT kill the CLI — let the agent finish its work
+      // Only clean up WS-related mappings
+      for (const sessionId of streams) {
+        this.sessionToWs.delete(sessionId);
+
+        if (this.permissionService) {
+          this.permissionService.cancelPendingRequests(sessionId);
+        }
+
+        this.persistPartialMessages(sessionId);
+
+        this.log('CHAT_WS', `WS disconnected for session ${sessionId} — CLI continues running`);
+      }
+      this.activeStreams.delete(ws);
+    }
+  };
+
+  /**
+   * Handle attach_session: client reconnects and wants to re-attach to an in-flight stream.
+   * If an active stream exists, re-maps the WS and sends catch-up content.
+   */
+  private handleAttachSession = (ws: WebSocket, msg: AttachSessionPayload): void => {
+    const { claudeSessionId } = msg;
+    const ctx = this.streamContexts.get(claudeSessionId);
+
+    if (!ctx) {
+      // No active stream — the CLI has already finished. Nothing to attach to.
+      return;
+    }
+
+    // Re-map the new WS to receive events for this session
+    this.sessionToWs.set(claudeSessionId, ws);
+    let activeSet = this.activeStreams.get(ws);
+    if (!activeSet) {
+      activeSet = new Set();
+      this.activeStreams.set(ws, activeSet);
+    }
+    activeSet.add(claudeSessionId);
+
+    // Send stream_resumed to signal re-attachment
+    this.sendWsMessage(ws, { type: 'stream_resumed', claudeSessionId });
+
+    // Send accumulated assistant content as a catch-up stream_event
+    if (ctx.lastAssistantContent !== '[]') {
+      const catchUpEvent: Record<string, unknown> = {
+        type: 'assistant',
+        message: { content: JSON.parse(ctx.lastAssistantContent) },
+      };
+      this.sendWsMessage(ws, { type: 'stream_event', event: catchUpEvent });
+    }
+
+    this.log('CHAT_WS', `Client re-attached to in-flight stream for session ${claudeSessionId}`);
+  };
+}