TaskStateMachine Deep Dive Analysis

interface TaskStateData {
  taskId: number                // Task identifier
  status: TaskStatus           // Current state machine status
  messages: Map<string, UnifiedMessage>  // All messages
  streamingSubtaskId: number | null      // Active streaming subtask
  streamingInfo: StreamingInfo | null    // Redis cached content
  error: string | null                   // Error state
  isStopping: boolean                    // Stop requested
}

type TaskStatus = 
  | 'idle'        // Not joined WebSocket room
  | 'joining'     // Joining WebSocket room
  | 'syncing'     // Syncing messages from backend
  | 'ready'       // Ready, no streaming
  | 'streaming'   // Has active streaming message
  | 'error'       // Error occurred

// In TaskStateMachine
private listeners: Set<StateListener> = new Set()

subscribe(listener: StateListener): () => void {
  this.listeners.add(listener)
  return () => { this.listeners.delete(listener) }
}

private notifyListeners(): void {
  const stateCopy = this.getState()
  this.listeners.forEach(listener => {
    try {
      listener(stateCopy)
    } catch (err) {
      console.error('[TaskStateMachine] Error in listener:', err)
    }
  })
}

private async handleRecover(event: Extract<Event, { type: 'RECOVER' }>): Promise<void> {
  // 1. Queue if already joining/syncing
  if (this.state.status === 'joining' || this.state.status === 'syncing') {
    this.pendingRecovery = true
    return
  }

  // 2. Debounce check (1s default)
  const now = Date.now()
  if (!event.force && now - this.lastRecoveryTime < this.recoveryDebounceMs) {
    return
  }
  this.lastRecoveryTime = now

  // 3. Check WebSocket connection
  if (!this.deps.isConnected()) return

  // 4. Calculate max messageId for incremental sync
  let maxMessageId: number | undefined
  for (const msg of this.state.messages.values()) {
    if (msg.messageId !== undefined && 
        (maxMessageId === undefined || msg.messageId > maxMessageId)) {
      maxMessageId = msg.messageId
    }
  }

  // 5. Join task room with incremental sync hint
  const response = await this.deps.joinTask(this.state.taskId, {
    forceRefresh: true,
    afterMessageId: maxMessageId,  // Only get messages after this ID
  })

  // 6. Dispatch JOIN_SUCCESS with subtasks for sync
  await this.dispatch({
    type: 'JOIN_SUCCESS',
    streamingInfo: response.streaming,
    subtasks: response.subtasks as TaskDetailSubtask[] | undefined,
  })
}

// Content priority in buildMessages()
if (!isUserMessage && subtask.status === 'RUNNING') {
  let bestContent = backendContent
  
  // Check existing message content
  if (existingAiMessage && existingAiMessage.content.length > bestContent.length) {
    bestContent = existingAiMessage.content
  }
  
  // Check Redis cached content (highest priority)
  if (streamingInfo && 
      streamingInfo.subtask_id === subtask.id &&
      streamingInfo.cached_content &&
      streamingInfo.cached_content.length > bestContent.length) {
    bestContent = streamingInfo.cached_content
  }
  
  messages.set(messageId, { ...bestContent })
}

private handleChatChunkEvent(event: Extract<Event, { type: 'CHAT_CHUNK' }>): void {
  const aiMessageId = generateMessageId('ai', event.subtaskId)

  // CRITICAL: Queue chunks during joining/syncing
  if (this.state.status === 'joining' || this.state.status === 'syncing') {
    this.pendingChunks.push({
      subtaskId: event.subtaskId,
      content: event.content,
      result: event.result,
      sources: event.sources,
      blockId: event.blockId,
    })
    return
  }

  const existingMessage = this.getAiMessage(event.subtaskId)
  if (!existingMessage) {
    console.warn('[TaskStateMachine] CHAT_CHUNK ignored - message not found')
    return
  }

  // CRITICAL FIX: Always append content to message.content
  // Previously, when block_id existed, content was only added to blocks
  // causing UI to lose cached_content after page refresh
  const updatedMessage: UnifiedMessage = {
    ...existingMessage,
    content: existingMessage.content + event.content,
  }

  // Handle reasoning content (DeepSeek R1)
  if (event.result?.reasoning_chunk) {
    updatedMessage.reasoningContent =
      (existingMessage.reasoningContent || '') + event.result.reasoning_chunk
  }

  // Handle blocks (text blocks, tool calls)
  const mergedBlocks = this.mergeBlocks(existingMessage, event)
  if (event.result) {
    updatedMessage.result = {
      ...existingMessage.result,
      ...event.result,
      blocks: mergedBlocks,
    }
  }

  // Update state and notify
  const newMessages = new Map(this.state.messages)
  newMessages.set(aiMessageId, updatedMessage)
  this.state = { ...this.state, messages: newMessages }
  this.notifyListeners()
}

private pendingChunks: PendingChunkEvent[] = []

// Queue chunks during sync
if (this.state.status === 'joining' || this.state.status === 'syncing') {
  this.pendingChunks.push({ subtaskId, content, result, sources, blockId })
  return
}

// Apply after sync completes
private applyPendingChunks(): void {
  for (const chunk of this.pendingChunks) {
    const aiMessageId = generateMessageId('ai', chunk.subtaskId)
    const existingMessage = this.state.messages.get(aiMessageId)
    
    if (!existingMessage) {
      console.warn('[TaskStateMachine] Pending chunk skipped - message not found')
      continue
    }
    
    // Apply chunk to message (same logic as handleChatChunkEvent)
    const updatedMessage = {
      ...existingMessage,
      content: existingMessage.content + chunk.content,
    }
    // ... handle blocks, reasoning, sources
  }
  
  this.pendingChunks = []
  this.notifyListeners()
}

migrateState(tempTaskId: number, realTaskId: number): void {
  const tempMachine = this.machines.get(tempTaskId)
  if (!tempMachine) {
    console.log('[TaskStateManager] No temp machine to migrate from', tempTaskId)
    return
  }

  // Check if real machine already exists (edge case)
  const existingRealMachine = this.machines.get(realTaskId)
  if (existingRealMachine && existingRealMachine !== tempMachine) {
    console.warn('[TaskStateManager] Real machine already exists')
    tempMachine.leave()
    this.machines.delete(tempTaskId)
    return
  }

  // KEY: Make real task ID point to SAME state machine instance
  // IMPORTANT: Keep BOTH mappings so components using either ID work
  this.machines.set(realTaskId, tempMachine)
  // DO NOT: this.machines.delete(tempTaskId)  // REMOVED!

  const state = tempMachine.getState()
  console.log('[TaskStateManager] Migrated state: both', tempTaskId, 
              'and', realTaskId, 'now point to same machine')
}

// CRITICAL: If streamingInfo exists but no streaming message was created,
// create one now. This handles the race condition where:
// 1. User joins room while streaming is in progress
// 2. Backend returns streamingInfo with subtask_id and cached_content
// 3. But subtasks array doesn't contain this subtask yet
const streamingInfo = this.state.streamingInfo
if (streamingInfo && streamingInfo.subtask_id) {
  const aiMessageId = generateMessageId('ai', streamingInfo.subtask_id)
  const existingMessage = this.state.messages.get(aiMessageId)

  if (!existingMessage) {
    const newMessages = new Map(this.state.messages)
    newMessages.set(aiMessageId, {
      id: aiMessageId,
      type: 'ai',
      status: 'streaming',
      content: streamingInfo.cached_content || '',  // Use Redis cache!
      timestamp: Date.now(),
      subtaskId: streamingInfo.subtask_id,
    })
    this.state = { ...this.state, messages: newMessages }
  }
}

// Sort messages by messageId (primary) and timestamp (secondary)
const sortedMessages = messages.sort((a, b) => {
  // If both have messageId, use it as primary sort key
  if (a.messageId !== undefined && b.messageId !== undefined) {
    if (a.messageId !== b.messageId) {
      return a.messageId - b.messageId
    }
    // Same messageId, use timestamp as secondary sort key
    return a.timestamp - b.timestamp
  }
  // If only one has messageId, the one with messageId comes first
  if (a.messageId !== undefined) return -1
  if (b.messageId !== undefined) return 1
  // Neither has messageId (both pending), sort by timestamp
  return a.timestamp - b.timestamp
})

// ❌ WRONG
const messages = selectedTaskDetail.subtasks  // Stale data!

// ✅ CORRECT
const { messages } = useUnifiedMessages({ team, isGroupChat })

// WebSocket events update state machine:
handleChatStart(data) {
  const machine = taskStateManager.getOrCreate(data.task_id)
  machine.handleChatStart(data.subtask_id)  // Updates internal state
}

handleChatChunk(data) {
  const machine = taskStateManager.getOrCreate(data.task_id)
  machine.handleChatChunk(data.subtask_id, data.content)
}

// Recovery fetches subtasks and builds messages
private async doSync(subtasks?: TaskDetailSubtask[]): Promise<void> {
  if (subtasks && subtasks.length > 0) {
    this.buildMessages(subtasks)  // Sync into state machine
  }
}

interface MessageBlock {
  id: string                    // Unique block identifier
  type: 'text' | 'tool' | 'thinking' | 'error'
  content?: string              // Text content
  tool_use_id?: string          // Tool call ID
  tool_name?: string            // Tool name
  tool_input?: Record<string, unknown>
  tool_output?: unknown
  status?: 'pending' | 'streaming' | 'done' | 'error'
  timestamp?: number
}

private mergeBlocks(
  existingMessage: UnifiedMessage,
  event: Extract<Event, { type: 'CHAT_CHUNK' }>
): MessageBlock[] {
  const existingBlocks = existingMessage.result?.blocks || []

  // Case 1: block_id with content - text block update
  if (event.blockId && event.content) {
    const blocksMap = new Map(existingBlocks.map(b => [b.id, b]))
    const targetBlock = blocksMap.get(event.blockId)

    if (targetBlock && targetBlock.type === 'text') {
      // Update existing text block
      const updatedBlock = {
        ...targetBlock,
        content: (targetBlock.content || '') + event.content,
      }
      blocksMap.set(event.blockId, updatedBlock)
    } else if (!targetBlock) {
      // Create new text block
      const newBlock: MessageBlock = {
        id: event.blockId,
        type: 'text',
        content: event.content,
        status: 'streaming',
        timestamp: Date.now(),
      }
      blocksMap.set(event.blockId, newBlock)
    }

    return Array.from(blocksMap.values())
  }

  // Case 2: Tool blocks - merge by block.id (replace strategy)
  const incomingBlocks = event.result?.blocks || []
  if (incomingBlocks.length > 0) {
    const blocksMap = new Map(existingBlocks.map(b => [b.id, b]))
    incomingBlocks.forEach(incomingBlock => {
      blocksMap.set(incomingBlock.id, incomingBlock)  // Replace
    })
    return Array.from(blocksMap.values())
  }

  // Case 3: No changes
  return existingBlocks
}

// Per-task cleanup
leave(): void {
  this.state = {
    ...this.state,
    status: 'idle',
    messages: new Map(),           // Clear messages
    streamingSubtaskId: null,
    streamingInfo: null,
    error: null,
    isStopping: false,
  }
  this.pendingRecovery = false
  this.pendingChunks = []          // Clear pending chunks
}

// Cleanup via TaskStateManager
cleanup(taskId: number): void {
  const machine = this.machines.get(taskId)
  if (machine) {
    machine.leave()                  // Reset state
    this.machines.delete(taskId)     // Remove reference
  }
}

// Cleanup all
cleanupAll(): void {
  for (const [taskId, machine] of this.machines) {
    machine.leave()
  }
  this.machines.clear()
}