PicoBot/web/src/hooks/useChat.test.ts

import { renderHook, act } from '@testing-library/react';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { useChat } from './useChat';
import type {
  WsInbound,
  SessionEstablished,
  SessionList,
  SessionSummary,
  TopicList,
  TopicSummary,
  StreamDelta,
  AssistantResponse,
  ToolCall,
  ToolResult,
  ToolPending,
  WsError,
  TaskStarted,
  TaskMessagesLoaded,
  MemoryList,
  MemorySummary,
  SkillList,
  SkillSummary,
  TodoList,
  TodoItemSummary,
  ChannelList,
  Channel,
  SchedulerJobList,
  SchedulerJobSummary,
  SchedulerJobSessionLookup,
  ExecutionCancelled,
} from '../types/protocol';

// ---- helpers ----

function renderUseChat() {
  const sendMessage = vi.fn((_msg: WsInbound) => true);
  const { result } = renderHook(() => useChat());
  act(() => {
    result.current.setSendMessage(sendMessage);
  });
  return { result, sendMessage };
}

/** 取出 sendMessage 收到的最后一条 command payload（已 JSON.parse） */
function lastCommand(sendMessage: ReturnType<typeof vi.fn>): unknown {
  const calls = sendMessage.mock.calls;
  const last = calls.length > 0 ? (calls[calls.length - 1][0] as WsInbound) : undefined;
  if (last && last.type === 'command') {
    return JSON.parse(last.payload);
  }
  return undefined;
}

// ---- fixtures ----

const sessionEstablished: SessionEstablished = {
  type: 'session_established',
  session_id: 'sess-1',
};

function makeSession(id: string): SessionSummary {
  return {
    session_id: id,
    title: `Session ${id}`,
    channel_name: 'websocket',
    chat_id: `chat-${id}`,
    message_count: 0,
    last_active_at: 1000,
  };
}

const sessionList: SessionList = {
  type: 'session_list',
  sessions: [makeSession('s1'), makeSession('s2')],
};

function makeTopicSummary(id: string, sessionId = 's1'): TopicSummary {
  return {
    topic_id: id,
    session_id: sessionId,
    title: `Topic ${id}`,
    message_count: 0,
    created_at: 1000,
    last_active_at: 2000,
  };
}

const topicList: TopicList = {
  type: 'topic_list',
  topics: [makeTopicSummary('t1'), makeTopicSummary('t2')],
  session_id: 's1',
};

const streamDelta1: StreamDelta = {
  type: 'stream_delta',
  id: 'm1',
  delta: 'Hello',
};
const streamDelta2: StreamDelta = {
  type: 'stream_delta',
  id: 'm1',
  delta: ' world',
};

const assistantResponse: AssistantResponse = {
  type: 'assistant_response',
  id: 'm1',
  content: 'Hello world',
  role: 'assistant',
};

const toolCall: ToolCall = {
  type: 'tool_call',
  id: 'tc1',
  tool_call_id: 'tc1',
  tool_name: 'calculator',
  arguments: { x: 1 },
  content: 'calling calculator',
  role: 'tool',
};

const toolResult: ToolResult = {
  type: 'tool_result',
  id: 'tr1',
  tool_call_id: 'tc1',
  tool_name: 'calculator',
  content: '42',
  role: 'tool',
};

const toolPending: ToolPending = {
  type: 'tool_pending',
  id: 'tp1',
  tool_call_id: 'tp1',
  tool_name: 'bash',
  content: 'waiting',
  resume_hint: 'resume later',
  role: 'tool',
};

const errorMsg: WsError = {
  type: 'error',
  code: 'ERR',
  message: 'something broke',
};

const executionCancelled: ExecutionCancelled = {
  type: 'execution_cancelled',
  message: 'stopped by user',
};

const memoryList: MemoryList = {
  type: 'memory_list',
  memories: [
    { id: 'mem1', namespace: 'ns', memory_key: 'k', content: 'c', created_at: 1, updated_at: 2 },
  ] as MemorySummary[],
};

const skillList: SkillList = {
  type: 'skill_list',
  skills: [{ name: 'skill1', description: 'd', source: 'builtin' }] as SkillSummary[],
};

const todoList: TodoList = {
  type: 'todo_list',
  todos: [
    {
      id: 'todo1',
      content: 'task',
      status: 'pending',
      priority: 'high',
      created_at: 1,
      updated_at: 2,
    },
  ] as TodoItemSummary[],
  scope_key: 'main',
};

const channelList: ChannelList = {
  type: 'channel_list',
  channels: [
    { id: 'websocket', name: 'WebSocket', isWritable: true },
    { id: 'cli', name: 'CLI', isWritable: false },
  ] as Channel[],
};

const schedulerJobList: SchedulerJobList = {
  type: 'scheduler_job_list',
  jobs: [
    {
      id: 'job1',
      kind: 'one_off',
      schedule: {},
      enabled: true,
      state: 'idle',
      run_count: 0,
      created_at: 1,
    } as SchedulerJobSummary,
  ],
};

// ---- tests ----

beforeEach(() => {
  vi.clearAllMocks();
});

describe('useChat - handleServerMessage characterization', () => {
  it('1. session_established sets connectionId and isConnected', () => {
    const { result } = renderUseChat();
    expect(result.current.isConnected).toBe(false);
    act(() => result.current.handleServerMessage(sessionEstablished));
    expect(result.current.connectionId).toBe('sess-1');
    expect(result.current.isConnected).toBe(true);
  });

  it('2. session_list fills sessions and auto-selects the first', () => {
    const { result } = renderUseChat();
    act(() => result.current.handleServerMessage(sessionList));
    expect(result.current.sessions).toHaveLength(2);
    expect(result.current.selectedSessionId).toBe('s1');
    expect(result.current.session?.session_id).toBe('s1');
  });

  it('3. topic_list maps topics; after createTopic it auto-focuses the first (newest)', () => {
    const { result } = renderUseChat();
    // establish session + topic list to set baseline
    act(() => result.current.handleServerMessage(sessionEstablished));
    act(() => result.current.handleServerMessage(sessionList));
    // first topic_list (without createTopic) sets topics but does NOT auto-select
    act(() => result.current.handleServerMessage(topicList));
    expect(result.current.topics).toHaveLength(2);
    expect(result.current.selectedTopic).toBeNull();

    // simulate createTopic flow: pendingNewTopicRef set true, then new topic_list arrives
    act(() => result.current.createTopic('new topic'));
    const newTopicList: TopicList = {
      type: 'topic_list',
      topics: [makeTopicSummary('t3'), makeTopicSummary('t1'), makeTopicSummary('t2')],
      session_id: 's1',
    };
    act(() => result.current.handleServerMessage(newTopicList));
    expect(result.current.selectedTopic).toBe('t3');
  });

  it('4. stream_delta creates a message then accumulates into it by id', () => {
    const { result } = renderUseChat();
    act(() => result.current.handleServerMessage(streamDelta1));
    expect(result.current.messages).toHaveLength(1);
    expect(result.current.messages[0].content).toBe('Hello');
    act(() => result.current.handleServerMessage(streamDelta2));
    expect(result.current.messages).toHaveLength(1);
    expect(result.current.messages[0].content).toBe('Hello world');
  });

  it('5. assistant_response replaces the streamed message by id', () => {
    const { result } = renderUseChat();
    act(() => result.current.handleServerMessage(streamDelta1));
    act(() => result.current.handleServerMessage(streamDelta2));
    act(() => result.current.handleServerMessage(assistantResponse));
    expect(result.current.messages).toHaveLength(1);
    expect(result.current.messages[0].content).toBe('Hello world');
    expect(result.current.messages[0].id).toBe('m1');
  });

  it('6. tool_call / tool_result / tool_pending append corresponding message types', () => {
    const { result } = renderUseChat();
    act(() => result.current.handleServerMessage(toolCall));
    act(() => result.current.handleServerMessage(toolResult));
    act(() => result.current.handleServerMessage(toolPending));
    expect(result.current.messages).toHaveLength(3);
    expect(result.current.messages[0].type).toBe('tool_call');
    expect(result.current.messages[0].toolName).toBe('calculator');
    expect(result.current.messages[1].type).toBe('tool_result');
    expect(result.current.messages[2].type).toBe('tool_pending');
    expect(result.current.messages[2].content).toContain('resume later');
  });

  it('7. error and execution_cancelled append a message and clear isLoading', () => {
    const { result } = renderUseChat();
    // set isLoading true via handleMessage
    act(() => result.current.handleMessage('hi'));
    expect(result.current.isLoading).toBe(true);
    act(() => result.current.handleServerMessage(errorMsg));
    expect(result.current.isLoading).toBe(false);
    const errMsg = result.current.messages[result.current.messages.length - 1];
    expect(errMsg?.content).toBe('Error: something broke');

    // reset isLoading + cleared, then test execution_cancelled
    act(() => result.current.handleMessage('hi again'));
    expect(result.current.isLoading).toBe(true);
    act(() => result.current.handleServerMessage(executionCancelled));
    expect(result.current.isLoading).toBe(false);
    const cancelMsg = result.current.messages[result.current.messages.length - 1];
    expect(cancelMsg?.content).toBe('stopped by user');
  });

  it('8. memory_list / skill_list / todo_list / channel_list / scheduler_job_list set corresponding state', () => {
    const { result } = renderUseChat();
    act(() => result.current.handleServerMessage(memoryList));
    act(() => result.current.handleServerMessage(skillList));
    act(() => result.current.handleServerMessage(todoList));
    act(() => result.current.handleServerMessage(channelList));
    act(() => result.current.handleServerMessage(schedulerJobList));
    expect(result.current.memories).toHaveLength(1);
    expect(result.current.skills).toHaveLength(1);
    expect(result.current.todos).toHaveLength(1);
    expect(result.current.channels).toHaveLength(2);
    expect(result.current.schedulerJobs).toHaveLength(1);
  });

  it('9. task_started (main view, no parent) backfills navigateToTaskId on matching task tool_call', () => {
    const { result } = renderUseChat();
    const taskToolCall: ToolCall = {
      type: 'tool_call',
      id: 'tc-task',
      tool_call_id: 'tc-task',
      tool_name: 'task',
      arguments: { prompt: 'do sub' },
      content: 'spawning sub',
      role: 'tool',
    };
    act(() => result.current.handleServerMessage(taskToolCall));
    expect(
      result.current.messages[result.current.messages.length - 1]?.navigateToTaskId,
    ).toBeUndefined();

    const taskStarted: TaskStarted = {
      type: 'task_started',
      task_id: 'sub-1',
      description: 'sub agent',
      subagent_type: 'general',
      tool_call_id: 'tc-task',
    };
    act(() => result.current.handleServerMessage(taskStarted));
    expect(result.current.messages[result.current.messages.length - 1]?.navigateToTaskId).toBe(
      'sub-1',
    );
  });

  it('10. sub-agent view: task_messages_loaded updates stack top; tagged messages route to sub view not main', () => {
    const { result } = renderUseChat();
    // enter sub-agent view for task "sub-1"
    act(() => result.current.enterSubAgentView('sub-1', 'sub agent', 'general'));
    expect(result.current.subAgentView?.taskId).toBe('sub-1');

    // task_messages_loaded updates top metadata
    const loaded: TaskMessagesLoaded = {
      type: 'task_messages_loaded',
      task_id: 'sub-1',
      description: 'sub agent',
      subagent_type: 'general',
      status: 'running',
      summary: 'working',
    };
    act(() => result.current.handleServerMessage(loaded));
    expect(result.current.subAgentView?.status).toBe('running');
    expect(result.current.subAgentView?.summary).toBe('working');

    // a stream_delta tagged with subagent_task_id === 'sub-1' goes to sub view, not main
    const subStream: StreamDelta = {
      type: 'stream_delta',
      id: 'sub-m1',
      delta: 'sub hello',
      subagent_task_id: 'sub-1',
    };
    act(() => result.current.handleServerMessage(subStream));
    expect(result.current.subAgentView?.messages).toHaveLength(1);
    expect(result.current.subAgentView?.messages[0].content).toBe('sub hello');

    // exit back to main: main messages should not contain the sub-agent message
    act(() => result.current.exitSubAgentView());
    expect(result.current.messages.find((m) => m.id === 'sub-m1')).toBeUndefined();
  });

  it('11. scheduler view: chat messages route into schedulerView.messages, not main', () => {
    const { result } = renderUseChat();
    const lookup: SchedulerJobSessionLookup = { channel: 'scheduler', chat_id: 'job-chat' };
    act(() => result.current.enterSchedulerJobView(lookup, 'job1', 'job desc'));
    expect(result.current.schedulerView).not.toBeNull();

    act(() => result.current.handleServerMessage(assistantResponse));
    expect(result.current.schedulerView?.messages).toHaveLength(1);
    expect(result.current.schedulerView?.messages[0].content).toBe('Hello world');

    // exit scheduler view: main messages should not contain the routed message
    act(() => result.current.exitSchedulerJobView());
    expect(result.current.messages.find((m) => m.id === 'm1')).toBeUndefined();
  });

  it('12. tool_result with tool_name=todo_write triggers a list_todos command in main view', () => {
    const { result, sendMessage } = renderUseChat();
    const todoWriteResult: ToolResult = {
      type: 'tool_result',
      id: 'tr-todo',
      tool_call_id: 'tc-todo',
      tool_name: 'todo_write',
      content: 'updated',
      role: 'tool',
    };
    act(() => result.current.handleServerMessage(todoWriteResult));
    const cmd = lastCommand(sendMessage);
    expect(cmd).toEqual({ type: 'list_todos' });
  });

  it('13. stream_delta whose topic_id does not match selectedTopic is discarded', () => {
    const { result } = renderUseChat();
    act(() => {
      result.current.handleServerMessage(sessionEstablished);
      result.current.handleServerMessage(sessionList);
      result.current.handleServerMessage(topicList);
    });
    // topic_list without createTopic does NOT auto-select; manually select t1
    act(() => result.current.selectTopic('t1'));
    expect(result.current.selectedTopic).toBe('t1');

    const otherTopicStream: StreamDelta = {
      type: 'stream_delta',
      id: 'm-other',
      delta: 'should be dropped',
      topic_id: 't-other',
    };
    act(() => result.current.handleServerMessage(otherTopicStream));
    expect(result.current.messages.find((m) => m.id === 'm-other')).toBeUndefined();
  });

  it('14. multiple stream_delta in one act() batch accumulate correctly', () => {
    const { result } = renderUseChat();
    // 模拟同一同步批次内连续到达多个 delta（验证 chunks 累加 + index 跟踪）
    act(() => {
      result.current.handleServerMessage({ type: 'stream_delta', id: 'batch1', delta: 'A' });
      result.current.handleServerMessage({ type: 'stream_delta', id: 'batch1', delta: 'B' });
      result.current.handleServerMessage({ type: 'stream_delta', id: 'batch1', delta: 'C' });
    });
    expect(result.current.messages).toHaveLength(1);
    expect(result.current.messages[0].id).toBe('batch1');
    expect(result.current.messages[0].content).toBe('ABC');
  });

  it('15. stream_end then assistant_response replaces streamed content including reasoning', () => {
    const { result } = renderUseChat();
    act(() => {
      result.current.handleServerMessage({
        type: 'stream_delta',
        id: 'r1',
        delta: 'text',
        reasoning_delta: 'think',
      });
    });
    expect(result.current.messages[0].content).toBe('text');
    expect(result.current.messages[0].reasoningContent).toBe('think');
    // stream_end 触发 finishStreaming（flush pending + clear）
    act(() => result.current.handleServerMessage({ type: 'stream_end', id: 'r1' }));
    // assistant_response 替换整条消息
    act(() =>
      result.current.handleServerMessage({
        type: 'assistant_response',
        id: 'r1',
        content: 'final text',
        role: 'assistant',
        reasoning_content: 'final reasoning',
      }),
    );
    expect(result.current.messages).toHaveLength(1);
    expect(result.current.messages[0].content).toBe('final text');
    expect(result.current.messages[0].reasoningContent).toBe('final reasoning');
  });
});