@assistkick/create 1.9.0 → 1.11.0

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

@@ -0,0 +1,1243 @@
+import { describe, it, mock, beforeEach } from 'node:test';
+import assert from 'node:assert/strict';
+import { ChatWsHandler, type SendMessagePayload, type ChatServerMessage } from './chat_ws_handler.ts';
+
+/**
+ * Create a mock WebSocket with controllable readyState and message capture.
+ */
+const createMockWs = (readyState = 1) => {
+  const messages: string[] = [];
+  const listeners: Record<string, Function[]> = {};
+
+  return {
+    OPEN: 1,
+    readyState,
+    send: mock.fn((data: string) => { messages.push(data); }),
+    close: mock.fn(),
+    on: mock.fn((event: string, cb: Function) => {
+      if (!listeners[event]) listeners[event] = [];
+      listeners[event].push(cb);
+    }),
+    _messages: messages,
+    _emit: (event: string, ...args: unknown[]) => {
+      for (const cb of listeners[event] || []) cb(...args);
+    },
+    _getParsedMessages: (): ChatServerMessage[] => messages.map(m => JSON.parse(m)),
+  };
+};
+
+/**
+ * Create a mock WebSocketServer.
+ */
+const createMockWss = () => ({
+  handleUpgrade: mock.fn((_req: unknown, _socket: unknown, _head: unknown, cb: Function) => {
+    cb(createMockWs());
+  }),
+});
+
+/**
+ * Create a mock AuthService.
+ */
+const createMockAuthService = (role = 'admin') => ({
+  verifyToken: mock.fn(async () => ({ sub: 'user-1', email: 'admin@test.com', role })),
+});
+
+/**
+ * Create a mock ChatCliBridge.
+ */
+const createMockBridge = () => {
+  let lastOnEvent: ((event: Record<string, unknown>) => void) | null = null;
+  let resolveCompletion: ((result: { exitCode: number | null; claudeSessionId: string; sessionId: string | null }) => void) | null = null;
+  let rejectCompletion: ((err: Error) => void) | null = null;
+
+  return {
+    spawn: mock.fn((options: Record<string, unknown>) => {
+      lastOnEvent = options.onEvent as (event: Record<string, unknown>) => void;
+      const completion = new Promise<{ exitCode: number | null; claudeSessionId: string; sessionId: string | null }>((resolve, reject) => {
+        resolveCompletion = resolve;
+        rejectCompletion = reject;
+      });
+      return { childProcess: {}, completion, systemPrompt: options.isNewSession ? 'test system prompt' : null };
+    }),
+    kill: mock.fn(() => true),
+    isActive: mock.fn(() => false),
+    _emitEvent: (event: Record<string, unknown>) => {
+      if (lastOnEvent) lastOnEvent(event);
+    },
+    _completeStream: (exitCode: number | null, claudeSessionId: string) => {
+      if (resolveCompletion) resolveCompletion({ exitCode, claudeSessionId, sessionId: null });
+    },
+    _failStream: (err: Error) => {
+      if (rejectCompletion) rejectCompletion(err);
+    },
+  };
+};
+
+/**
+ * Create a mock ChatMessageRepository for testing persistence.
+ */
+const createMockMessageRepository = () => {
+  let msgCounter = 0;
+  return {
+    getNextOrderIndex: mock.fn(async () => 0),
+    saveMessage: mock.fn(async (sessionId: string, role: string, content: string, orderIndex: number) => ({
+      id: `cmsg_mock_${++msgCounter}`,
+      sessionId,
+      role,
+      content,
+      orderIndex,
+      createdAt: new Date().toISOString(),
+    })),
+    updateMessageContent: mock.fn(async () => {}),
+    getMessages: mock.fn(async () => []),
+    deleteBySessionId: mock.fn(async () => 0),
+  };
+};
+
+/**
+ * Create a mock IncomingMessage with cookie headers.
+ */
+const createMockRequest = (path: string, token?: string) => ({
+  url: path,
+  headers: {
+    cookie: token ? `access_token=${token}` : undefined,
+  },
+});
+
+describe('ChatWsHandler', () => {
+  let handler: ChatWsHandler;
+  let mockWss: ReturnType<typeof createMockWss>;
+  let mockAuthService: ReturnType<typeof createMockAuthService>;
+  let mockBridge: ReturnType<typeof createMockBridge>;
+  let logMock: ReturnType<typeof mock.fn>;
+
+  beforeEach(() => {
+    logMock = mock.fn();
+    mockWss = createMockWss();
+    mockAuthService = createMockAuthService();
+    mockBridge = createMockBridge();
+
+    handler = new ChatWsHandler({
+      wss: mockWss as any,
+      authService: mockAuthService as any,
+      chatCliBridge: mockBridge as any,
+      log: logMock,
+    });
+  });
+
+  describe('handleUpgrade', () => {
+    it('calls wss.handleUpgrade for /api/chat path', () => {
+      const req = createMockRequest('/api/chat', 'valid-token');
+      const socket = {} as any;
+      const head = Buffer.from('');
+
+      handler.handleUpgrade(req as any, socket, head);
+
+      assert.equal(mockWss.handleUpgrade.mock.calls.length, 1);
+    });
+
+    it('ignores non-chat paths', () => {
+      const req = createMockRequest('/api/terminal', 'valid-token');
+      const socket = {} as any;
+      const head = Buffer.from('');
+
+      handler.handleUpgrade(req as any, socket, head);
+
+      assert.equal(mockWss.handleUpgrade.mock.calls.length, 0);
+    });
+
+    it('ignores other paths like /api/other', () => {
+      const req = createMockRequest('/api/other', 'valid-token');
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      assert.equal(mockWss.handleUpgrade.mock.calls.length, 0);
+    });
+  });
+
+  describe('authentication', () => {
+    it('rejects connection with no access token', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+
+      // Allow async onConnection to complete
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      assert.equal(ws.close.mock.calls.length, 1);
+      assert.equal(ws.close.mock.calls[0].arguments[0], 4001);
+    });
+
+    it('rejects connection with invalid token', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'bad-token');
+
+      const failAuthService = {
+        verifyToken: mock.fn(async () => { throw new Error('Invalid'); }),
+      };
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: failAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      assert.equal(ws.close.mock.calls.length, 1);
+      assert.equal(ws.close.mock.calls[0].arguments[0], 4001);
+    });
+
+    it('rejects non-admin users', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      const nonAdminAuth = createMockAuthService('viewer');
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: nonAdminAuth as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      assert.equal(ws.close.mock.calls.length, 1);
+      assert.equal(ws.close.mock.calls[0].arguments[0], 4003);
+    });
+
+    it('accepts admin users and registers message listener', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      assert.equal(ws.close.mock.calls.length, 0);
+      // Should have registered 'message', 'close', 'error' listeners
+      const registeredEvents = ws.on.mock.calls.map((c: any) => c.arguments[0]);
+      assert.ok(registeredEvents.includes('message'));
+      assert.ok(registeredEvents.includes('close'));
+      assert.ok(registeredEvents.includes('error'));
+    });
+  });
+
+  describe('send_message handling', () => {
+    let ws: ReturnType<typeof createMockWs>;
+
+    beforeEach(async () => {
+      ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+    });
+
+    it('spawns CLI and sends stream_start on send_message', () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello Claude',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+
+      // Should have called bridge.spawn
+      assert.equal(mockBridge.spawn.mock.calls.length, 1);
+      const spawnArgs = mockBridge.spawn.mock.calls[0].arguments[0];
+      assert.equal(spawnArgs.message, 'Hello Claude');
+      assert.equal(spawnArgs.projectId, 'proj_test');
+      assert.equal(spawnArgs.claudeSessionId, 'session-1');
+      assert.equal(spawnArgs.isNewSession, true);
+      assert.equal(spawnArgs.permissionMode, 'skip');
+
+      // Should have sent stream_start
+      const messages = ws._getParsedMessages();
+      assert.equal(messages[0].type, 'stream_start');
+      assert.equal((messages[0] as any).claudeSessionId, 'session-1');
+    });
+
+    it('forwards stream events to the WebSocket', () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+
+      // Simulate CLI emitting an assistant event
+      const assistantEvent = {
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'Hi there!' }] },
+      };
+      mockBridge._emitEvent(assistantEvent);
+
+      const messages = ws._getParsedMessages();
+      // stream_start + system_prompt + stream_event
+      assert.equal(messages.length, 3);
+      assert.equal(messages[2].type, 'stream_event');
+      assert.deepEqual((messages[2] as any).event, assistantEvent);
+    });
+
+    it('sends stream_end on successful completion', async () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+      mockBridge._completeStream(0, 'session-1');
+
+      // Wait for promise resolution
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      const messages = ws._getParsedMessages();
+      const endMsg = messages.find((m: any) => m.type === 'stream_end');
+      assert.ok(endMsg);
+      assert.equal((endMsg as any).exitCode, 0);
+      assert.equal((endMsg as any).claudeSessionId, 'session-1');
+    });
+
+    it('sends stream_error on CLI failure', async () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+      mockBridge._failStream(new Error('CLI crashed'));
+
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      const messages = ws._getParsedMessages();
+      const errMsg = messages.find((m: any) => m.type === 'stream_error');
+      assert.ok(errMsg);
+      assert.equal((errMsg as any).error, 'CLI crashed');
+    });
+
+    it('rejects concurrent messages on the same claude session', async () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'First message',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+
+      // Send second message to the SAME session before first completes
+      const msg2: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Second message',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: false,
+      };
+
+      ws._emit('message', JSON.stringify(msg2));
+
+      // Should only have spawned once
+      assert.equal(mockBridge.spawn.mock.calls.length, 1);
+
+      // Should have sent error for second message
+      const messages = ws._getParsedMessages();
+      const errMsg = messages.find((m: any) => m.type === 'error');
+      assert.ok(errMsg);
+      assert.ok((errMsg as any).message.includes('already being processed'));
+    });
+
+    it('allows concurrent messages to different claude sessions on the same connection', async () => {
+      const msg1: SendMessagePayload = {
+        type: 'send_message',
+        message: 'First session message',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg1));
+
+      // Send message to a DIFFERENT session while first is still streaming
+      const msg2: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Second session message',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-2',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg2));
+
+      // Should have spawned twice — one per session
+      assert.equal(mockBridge.spawn.mock.calls.length, 2);
+      assert.equal(mockBridge.spawn.mock.calls[0].arguments[0].claudeSessionId, 'session-1');
+      assert.equal(mockBridge.spawn.mock.calls[1].arguments[0].claudeSessionId, 'session-2');
+
+      // Should have sent two stream_start messages
+      const messages = ws._getParsedMessages();
+      const starts = messages.filter((m: any) => m.type === 'stream_start');
+      assert.equal(starts.length, 2);
+
+      // No error messages
+      const errors = messages.filter((m: any) => m.type === 'error');
+      assert.equal(errors.length, 0);
+    });
+
+    it('allows new message after stream completes', async () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'First',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+      mockBridge._completeStream(0, 'session-1');
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Now send second message — should work
+      const msg2: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Second',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: false,
+      };
+
+      ws._emit('message', JSON.stringify(msg2));
+
+      assert.equal(mockBridge.spawn.mock.calls.length, 2);
+    });
+
+    it('rejects messages with missing required fields', () => {
+      const msg = {
+        type: 'send_message',
+        message: '',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+
+      assert.equal(mockBridge.spawn.mock.calls.length, 0);
+      const messages = ws._getParsedMessages();
+      const errMsg = messages.find((m: any) => m.type === 'error');
+      assert.ok(errMsg);
+      assert.ok((errMsg as any).message.includes('Missing required fields'));
+    });
+
+    it('sends error on invalid JSON', () => {
+      ws._emit('message', 'not valid json');
+
+      const messages = ws._getParsedMessages();
+      assert.equal(messages[0].type, 'error');
+      assert.equal((messages[0] as any).message, 'Invalid JSON');
+    });
+
+    it('sends system_prompt event for new sessions', () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello Claude',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+
+      const messages = ws._getParsedMessages();
+      // stream_start + system_prompt
+      const systemPromptMsg = messages.find((m: any) => m.type === 'system_prompt');
+      assert.ok(systemPromptMsg, 'Expected a system_prompt message for new session');
+      assert.equal((systemPromptMsg as any).prompt, 'test system prompt');
+    });
+
+    it('does not send system_prompt event for resumed sessions', () => {
+      // First send a new session message to create the stream
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'First',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Complete the stream
+      mockBridge._completeStream(0, 'session-1');
+
+      // Clear messages
+      ws._messages.length = 0;
+
+      // Send resumed message — mock bridge returns systemPrompt: null for isNewSession: false
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Second',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: false,
+      }));
+
+      const messages = ws._getParsedMessages();
+      const systemPromptMsg = messages.find((m: any) => m.type === 'system_prompt');
+      assert.equal(systemPromptMsg, undefined, 'Should not send system_prompt for resumed sessions');
+    });
+
+    it('passes custom permission mode from message', () => {
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+        permissionMode: 'allowed_tools',
+        allowedTools: ['Read', 'Glob'],
+      };
+
+      ws._emit('message', JSON.stringify(msg));
+
+      const spawnArgs = mockBridge.spawn.mock.calls[0].arguments[0];
+      assert.equal(spawnArgs.permissionMode, 'allowed_tools');
+      assert.deepEqual(spawnArgs.allowedTools, ['Read', 'Glob']);
+    });
+  });
+
+  describe('connection cleanup', () => {
+    it('does NOT kill CLI process when WebSocket closes — lets agent continue', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+      ws._emit('message', JSON.stringify(msg));
+
+      // Close the WebSocket
+      ws._emit('close');
+
+      // CLI should NOT be killed — the agent continues running
+      assert.equal(mockBridge.kill.mock.calls.length, 0);
+    });
+
+    it('does NOT kill CLI process on WebSocket error — lets agent continue', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream
+      const msg: SendMessagePayload = {
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      };
+      ws._emit('message', JSON.stringify(msg));
+
+      // Trigger error
+      ws._emit('error', new Error('Connection lost'));
+
+      // CLI should NOT be killed — the agent continues running
+      assert.equal(mockBridge.kill.mock.calls.length, 0);
+    });
+
+    it('does not call kill when no active stream on close', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      ws._emit('close');
+
+      assert.equal(mockBridge.kill.mock.calls.length, 0);
+    });
+
+    it('persists partial messages when WebSocket disconnects mid-stream', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      const mockRepo = createMockMessageRepository();
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        chatMessageRepository: mockRepo as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream with a sessionId
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Build something',
+        projectId: 'proj_test',
+        sessionId: 'csess_abc',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Emit some assistant content
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'Working on it...' }] },
+      });
+
+      // Close the WebSocket (simulates page refresh)
+      ws._emit('close');
+
+      await new Promise(resolve => setTimeout(resolve, 20));
+
+      // Should have persisted partial messages
+      assert.equal(mockRepo.getNextOrderIndex.mock.calls.length, 1);
+      assert.equal(mockRepo.saveMessage.mock.calls.length, 2);
+
+      // First call: user message
+      assert.equal(mockRepo.saveMessage.mock.calls[0].arguments[0], 'csess_abc');
+      assert.equal(mockRepo.saveMessage.mock.calls[0].arguments[1], 'user');
+
+      // Second call: assistant partial response
+      assert.equal(mockRepo.saveMessage.mock.calls[1].arguments[0], 'csess_abc');
+      assert.equal(mockRepo.saveMessage.mock.calls[1].arguments[1], 'assistant');
+    });
+
+    it('updates assistant message with final content when CLI finishes after disconnect', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      const mockRepo = createMockMessageRepository();
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        chatMessageRepository: mockRepo as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Build something',
+        projectId: 'proj_test',
+        sessionId: 'csess_abc',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Emit partial content
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'Partial...' }] },
+      });
+
+      // WS disconnects
+      ws._emit('close');
+      await new Promise(resolve => setTimeout(resolve, 20));
+
+      // More content arrives after disconnect
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'Full response complete!' }] },
+      });
+
+      // CLI finishes
+      mockBridge._completeStream(0, 'session-1');
+      await new Promise(resolve => setTimeout(resolve, 20));
+
+      // Should have called updateMessageContent with the final content
+      assert.equal(mockRepo.updateMessageContent.mock.calls.length, 1);
+      // The assistant message is the 2nd saved (after user message), so it gets cmsg_mock_2
+      assert.equal(mockRepo.updateMessageContent.mock.calls[0].arguments[0], 'cmsg_mock_2');
+    });
+  });
+
+  describe('attach_session handling', () => {
+    it('re-attaches to in-flight stream and sends catch-up content', async () => {
+      const ws1 = createMockWs();
+      const req1 = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws1));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req1 as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream
+      ws1._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Emit some content
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'Working on it...' }] },
+      });
+
+      // WS disconnects
+      ws1._emit('close');
+
+      // New WS connects and attaches
+      const ws2 = createMockWs();
+      const req2 = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws2));
+      handler.handleUpgrade(req2 as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      ws2._emit('message', JSON.stringify({
+        type: 'attach_session',
+        claudeSessionId: 'session-1',
+      }));
+
+      const messages = ws2._getParsedMessages();
+      // Should have received: stream_resumed + stream_event (catch-up)
+      const resumed = messages.find((m: any) => m.type === 'stream_resumed');
+      assert.ok(resumed, 'Expected stream_resumed message');
+      assert.equal((resumed as any).claudeSessionId, 'session-1');
+
+      const catchUp = messages.find((m: any) => m.type === 'stream_event');
+      assert.ok(catchUp, 'Expected catch-up stream_event');
+      assert.equal((catchUp as any).event.type, 'assistant');
+    });
+
+    it('forwards new events to re-attached WS', async () => {
+      const ws1 = createMockWs();
+      const req1 = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws1));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req1 as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream
+      ws1._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // WS disconnects
+      ws1._emit('close');
+
+      // New WS connects and attaches
+      const ws2 = createMockWs();
+      const req2 = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws2));
+      handler.handleUpgrade(req2 as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      ws2._emit('message', JSON.stringify({
+        type: 'attach_session',
+        claudeSessionId: 'session-1',
+      }));
+
+      // Clear messages from re-attachment
+      ws2._messages.length = 0;
+
+      // New event arrives from CLI
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'New content after re-attach' }] },
+      });
+
+      const messages = ws2._getParsedMessages();
+      assert.equal(messages.length, 1);
+      assert.equal(messages[0].type, 'stream_event');
+      assert.equal((messages[0] as any).event.message.content[0].text, 'New content after re-attach');
+    });
+
+    it('does nothing when attaching to a session with no active stream', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      ws._emit('message', JSON.stringify({
+        type: 'attach_session',
+        claudeSessionId: 'nonexistent-session',
+      }));
+
+      // Should not send stream_resumed
+      const messages = ws._getParsedMessages();
+      const resumed = messages.find((m: any) => m.type === 'stream_resumed');
+      assert.equal(resumed, undefined);
+    });
+
+    it('sends stream_end to re-attached WS when CLI completes', async () => {
+      const ws1 = createMockWs();
+      const req1 = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws1));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req1 as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream
+      ws1._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // WS disconnects
+      ws1._emit('close');
+
+      // New WS connects and attaches
+      const ws2 = createMockWs();
+      const req2 = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws2));
+      handler.handleUpgrade(req2 as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      ws2._emit('message', JSON.stringify({
+        type: 'attach_session',
+        claudeSessionId: 'session-1',
+      }));
+
+      // CLI finishes
+      mockBridge._completeStream(0, 'session-1');
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Should have sent stream_end to ws2
+      const messages = ws2._getParsedMessages();
+      const endMsg = messages.find((m: any) => m.type === 'stream_end');
+      assert.ok(endMsg, 'Expected stream_end on re-attached WS');
+      assert.equal((endMsg as any).exitCode, 0);
+    });
+  });
+
+  describe('multiple stream events', () => {
+    it('forwards text, tool_use, and result events in order', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Build something',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Simulate streaming events from CLI
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'I will ' }] },
+      });
+
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'text', text: 'read the file.' }] },
+      });
+
+      mockBridge._emitEvent({
+        type: 'assistant',
+        message: { content: [{ type: 'tool_use', name: 'Read', input: { file_path: '/test.ts' } }] },
+      });
+
+      mockBridge._emitEvent({
+        type: 'result',
+        session_id: 'session-1',
+        cost_usd: 0.05,
+        duration_ms: 3000,
+      });
+
+      const messages = ws._getParsedMessages();
+      // stream_start + system_prompt + 4 stream_events
+      assert.equal(messages.length, 6);
+      assert.equal(messages[0].type, 'stream_start');
+      assert.equal(messages[1].type, 'system_prompt');
+      assert.equal(messages[2].type, 'stream_event');
+      assert.equal(messages[3].type, 'stream_event');
+      assert.equal(messages[4].type, 'stream_event');
+      assert.equal(messages[5].type, 'stream_event');
+
+      // Verify event contents
+      assert.equal((messages[2] as any).event.type, 'assistant');
+      assert.equal((messages[4] as any).event.type, 'assistant');
+      assert.equal((messages[4] as any).event.message.content[0].type, 'tool_use');
+      assert.equal((messages[5] as any).event.type, 'result');
+    });
+  });
+
+  describe('permission handling', () => {
+    it('forwards permission_response to permissionService', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+      const mockPermissionService = {
+        setPermissionRequestHandler: mock.fn(),
+        resolvePermission: mock.fn(async () => true),
+        cancelPendingRequests: mock.fn(),
+      };
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        permissionService: mockPermissionService as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Send a permission response
+      ws._emit('message', JSON.stringify({
+        type: 'permission_response',
+        requestId: 'perm_123',
+        projectId: 'proj_test',
+        decision: 'allow_once',
+      }));
+
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      assert.equal(mockPermissionService.resolvePermission.mock.calls.length, 1);
+      const resolveArgs = mockPermissionService.resolvePermission.mock.calls[0].arguments;
+      assert.equal(resolveArgs[0], 'perm_123');
+      assert.equal(resolveArgs[1], 'proj_test');
+      assert.equal(resolveArgs[2], 'allow_once');
+    });
+
+    it('cancels pending permission requests on WebSocket close', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+      const mockPermissionService = {
+        setPermissionRequestHandler: mock.fn(),
+        resolvePermission: mock.fn(async () => true),
+        cancelPendingRequests: mock.fn(),
+      };
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        permissionService: mockPermissionService as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream so there's an active session
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Close the WebSocket
+      ws._emit('close');
+
+      assert.equal(mockPermissionService.cancelPendingRequests.mock.calls.length, 1);
+      assert.equal(mockPermissionService.cancelPendingRequests.mock.calls[0].arguments[0], 'session-1');
+    });
+
+    it('works without permissionService (backward compatible)', async () => {
+      const ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Sending permission_response without service should not throw
+      ws._emit('message', JSON.stringify({
+        type: 'permission_response',
+        requestId: 'perm_123',
+        projectId: 'proj_test',
+        decision: 'allow_once',
+      }));
+
+      // Should not crash — just log a warning
+      assert.ok(true);
+    });
+  });
+
+  describe('stop_stream handling', () => {
+    let ws: ReturnType<typeof createMockWs>;
+
+    beforeEach(async () => {
+      ws = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+    });
+
+    it('kills CLI process and sends stream_cancelled on stop_stream', async () => {
+      // Start a stream
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Send stop_stream
+      ws._emit('message', JSON.stringify({ type: 'stop_stream' }));
+
+      // Should have called kill
+      assert.equal(mockBridge.kill.mock.calls.length, 1);
+      assert.equal(mockBridge.kill.mock.calls[0].arguments[0], 'session-1');
+
+      // Simulate CLI process exiting after kill
+      mockBridge._completeStream(null, 'session-1');
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Should have sent stream_cancelled (not stream_end)
+      const messages = ws._getParsedMessages();
+      const cancelled = messages.find((m: any) => m.type === 'stream_cancelled');
+      assert.ok(cancelled, 'Expected a stream_cancelled message');
+      assert.equal((cancelled as any).claudeSessionId, 'session-1');
+
+      // Should NOT have sent stream_end
+      const ended = messages.find((m: any) => m.type === 'stream_end');
+      assert.equal(ended, undefined, 'Should not send stream_end for cancelled stream');
+    });
+
+    it('sends error when no active stream to stop', () => {
+      ws._emit('message', JSON.stringify({ type: 'stop_stream' }));
+
+      const messages = ws._getParsedMessages();
+      const errMsg = messages.find((m: any) => m.type === 'error');
+      assert.ok(errMsg);
+      assert.ok((errMsg as any).message.includes('No active stream'));
+    });
+
+    it('allows new message after stop_stream completes', async () => {
+      // Start a stream
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: true,
+      }));
+
+      // Stop the stream
+      ws._emit('message', JSON.stringify({ type: 'stop_stream' }));
+      mockBridge._completeStream(null, 'session-1');
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Send another message — should work
+      ws._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Try again',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-1',
+        isNewSession: false,
+      }));
+
+      assert.equal(mockBridge.spawn.mock.calls.length, 2);
+    });
+
+    it('cancels pending permission requests on stop_stream', async () => {
+      const mockPermissionService = {
+        setPermissionRequestHandler: mock.fn(),
+        resolvePermission: mock.fn(async () => true),
+        cancelPendingRequests: mock.fn(),
+      };
+
+      const ws2 = createMockWs();
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws2));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        permissionService: mockPermissionService as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // Start a stream
+      ws2._emit('message', JSON.stringify({
+        type: 'send_message',
+        message: 'Hello',
+        projectId: 'proj_test',
+        claudeSessionId: 'session-2',
+        isNewSession: true,
+      }));
+
+      // Stop the stream
+      ws2._emit('message', JSON.stringify({ type: 'stop_stream' }));
+
+      assert.equal(mockPermissionService.cancelPendingRequests.mock.calls.length, 1);
+      assert.equal(mockPermissionService.cancelPendingRequests.mock.calls[0].arguments[0], 'session-2');
+    });
+  });
+
+  describe('does not send to closed WebSocket', () => {
+    it('skips sending when readyState is not OPEN', async () => {
+      const ws = createMockWs(3); // CLOSED
+      const req = createMockRequest('/api/chat', 'valid-token');
+
+      mockWss.handleUpgrade = mock.fn((_r: any, _s: any, _h: any, cb: Function) => cb(ws));
+      handler = new ChatWsHandler({
+        wss: mockWss as any,
+        authService: mockAuthService as any,
+        chatCliBridge: mockBridge as any,
+        log: logMock,
+      });
+
+      handler.handleUpgrade(req as any, {} as any, Buffer.from(''));
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      // No messages should be sent since readyState is not OPEN
+      assert.equal(ws._messages.length, 0);
+    });
+  });
+});