Claude-Code-Workflow/ccw/frontend/src/hooks/useWebSocket.test.tsx

// ========================================
// useWebSocket Hook Tests
// ========================================

import { act, renderHook } from '@testing-library/react';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { useWebSocket } from './useWebSocket';
import { useCliSessionStore } from '@/stores/cliSessionStore';
import { useExecutionMonitorStore } from '@/stores/executionMonitorStore';
import { useSessionManagerStore } from '@/stores/sessionManagerStore';
import { useWorkflowStore } from '@/stores/workflowStore';

class MockWebSocket {
  static readonly OPEN = 1;
  static instances: MockWebSocket[] = [];

  readonly url: string;
  readyState = 0;
  onopen: ((event: Event) => void) | null = null;
  onmessage: ((event: MessageEvent) => void) | null = null;
  onclose: ((event: CloseEvent) => void) | null = null;
  onerror: ((event: Event) => void) | null = null;
  send = vi.fn();
  close = vi.fn(() => {
    this.readyState = 3;
    this.onclose?.(new CloseEvent('close'));
  });

  constructor(url: string) {
    this.url = url;
    MockWebSocket.instances.push(this);
  }

  open() {
    this.readyState = MockWebSocket.OPEN;
    this.onopen?.(new Event('open'));
  }

  message(payload: unknown) {
    this.onmessage?.({ data: JSON.stringify(payload) } as MessageEvent);
  }
}

function createSession(sessionKey: string, workingDir = 'D:/workspace-a') {
  return {
    sessionKey,
    shellKind: 'pwsh',
    workingDir,
    createdAt: '2026-03-08T12:00:00.000Z',
    updatedAt: '2026-03-08T12:00:00.000Z',
    isPaused: false,
  };
}

function connectHook() {
  const hook = renderHook(() => useWebSocket({ enabled: true }));
  const socket = MockWebSocket.instances[MockWebSocket.instances.length - 1];
  if (!socket) {
    throw new Error('Expected WebSocket to be created');
  }

  act(() => {
    socket.open();
  });

  return { ...hook, socket };
}

describe('useWebSocket workspace scoping', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    localStorage.clear();
    MockWebSocket.instances = [];

    useCliSessionStore.getState().resetState();
    useExecutionMonitorStore.getState().resetState();
    useSessionManagerStore.getState().resetState();
    useWorkflowStore.setState({ projectPath: 'D:\\workspace-a' });

    vi.spyOn(console, 'log').mockImplementation(() => {});
    vi.stubGlobal('WebSocket', MockWebSocket as unknown as typeof WebSocket);
  });

  afterEach(() => {
    useCliSessionStore.getState().resetState();
    useExecutionMonitorStore.getState().resetState();
    useSessionManagerStore.getState().resetState();
    vi.restoreAllMocks();
    vi.unstubAllGlobals();
  });

  it('ignores scoped CLI and execution messages from another workspace', () => {
    const { socket } = connectHook();

    act(() => {
      socket.message({
        type: 'CLI_SESSION_CREATED',
        payload: {
          session: createSession('session-foreign', 'D:/workspace-b'),
          timestamp: '2026-03-08T12:00:01.000Z',
          projectPath: 'D:/workspace-b',
        },
      });
      socket.message({
        type: 'CLI_SESSION_LOCKED',
        payload: {
          sessionKey: 'session-foreign',
          reason: 'Foreign execution',
          executionId: 'exec-foreign',
          timestamp: '2026-03-08T12:00:02.000Z',
          projectPath: 'D:/workspace-b',
        },
      });
      socket.message({
        type: 'EXECUTION_STARTED',
        payload: {
          executionId: 'exec-foreign',
          flowId: 'flow-foreign',
          sessionKey: 'session-foreign',
          stepName: 'Foreign flow',
          totalSteps: 2,
          timestamp: '2026-03-08T12:00:03.000Z',
          projectPath: 'D:/workspace-b',
        },
      });
    });

    expect(useCliSessionStore.getState().sessions['session-foreign']).toBeUndefined();
    expect(useSessionManagerStore.getState().terminalMetas['session-foreign']).toBeUndefined();
    expect(useExecutionMonitorStore.getState().activeExecutions['exec-foreign']).toBeUndefined();
  });

  it('handles matching scoped messages and legacy messages for known sessions', () => {
    const { socket } = connectHook();

    act(() => {
      socket.message({
        type: 'CLI_SESSION_CREATED',
        payload: {
          session: createSession('session-local', 'D:/workspace-a/subdir'),
          timestamp: '2026-03-08T12:00:01.000Z',
          projectPath: 'd:/workspace-a',
        },
      });
      socket.message({
        type: 'CLI_SESSION_OUTPUT',
        payload: {
          sessionKey: 'session-local',
          data: 'hello from current workspace',
          timestamp: '2026-03-08T12:00:02.000Z',
        },
      });
      socket.message({
        type: 'CLI_SESSION_LOCKED',
        payload: {
          sessionKey: 'session-local',
          reason: 'Current execution',
          executionId: 'exec-local',
          timestamp: '2026-03-08T12:00:03.000Z',
          projectPath: 'D:/workspace-a',
        },
      });
      socket.message({
        type: 'EXECUTION_STARTED',
        payload: {
          executionId: 'exec-local',
          flowId: 'flow-local',
          sessionKey: 'session-local',
          stepName: 'Current flow',
          totalSteps: 3,
          timestamp: '2026-03-08T12:00:04.000Z',
        },
      });
    });

    const cliState = useCliSessionStore.getState();
    expect(cliState.sessions['session-local']?.workingDir).toBe('D:/workspace-a/subdir');
    expect(cliState.outputChunks['session-local']).toEqual([
      {
        data: 'hello from current workspace',
        timestamp: expect.any(Number),
      },
    ]);

    const sessionManagerState = useSessionManagerStore.getState();
    expect(sessionManagerState.terminalMetas['session-local']?.isLocked).toBe(true);
    expect(sessionManagerState.terminalMetas['session-local']?.lockedByExecutionId).toBe('exec-local');

    const executionState = useExecutionMonitorStore.getState();
    expect(executionState.activeExecutions['exec-local']?.sessionKey).toBe('session-local');
    expect(executionState.currentExecutionId).toBe('exec-local');
  });

  it('ignores legacy unscoped messages when session is unknown', () => {
    const { socket } = connectHook();

    act(() => {
      socket.message({
        type: 'CLI_SESSION_OUTPUT',
        payload: {
          sessionKey: 'session-unknown',
          data: 'should be ignored',
          timestamp: '2026-03-08T12:00:02.000Z',
        },
      });
      socket.message({
        type: 'EXECUTION_STARTED',
        payload: {
          executionId: 'exec-unknown',
          flowId: 'flow-unknown',
          sessionKey: 'session-unknown',
          stepName: 'Unknown flow',
          totalSteps: 1,
          timestamp: '2026-03-08T12:00:03.000Z',
        },
      });
    });

    expect(useCliSessionStore.getState().outputChunks['session-unknown']).toBeUndefined();
    expect(useExecutionMonitorStore.getState().activeExecutions['exec-unknown']).toBeUndefined();
  });
});