bookwiz.io / lib / hooks / useChatStreaming.ts
useChatStreaming.ts
Raw
import { useState, useRef, useCallback } from 'react'
import { ToolResult } from '@/lib/services/tool-executor'

// Removed StreamingActivity interface - now using React state directly

export interface UseChatStreamingReturn {
  // Streaming states
  isStreaming: boolean
  streamingMessage: string
  aiPlanning: string
  currentToolExecution: {
    name: string
    description: string
    args?: any
  } | null
  toolExecutionResults: Array<{
    name: string
    result: string
    success: boolean
  }>

  // Control functions
  startStreaming: () => void
  stopStreaming: () => void
  resetStreaming: () => void
  updateStreamingContent: (content: string) => void
  updateAiPlanning: (planning: string) => void
  updateCurrentExecution: (execution: {
    name: string
    description: string
    args?: any
  } | null) => void
  addToolResult: (result: {
    name: string
    result: string
    success: boolean
  }) => void

  // Abort controller
  abortControllerRef: React.MutableRefObject<AbortController | null>
  currentStreamingContentRef: React.MutableRefObject<string>

  // Utilities
  stopExecution: () => void
  processStreamChunk: (parsed: any, accumulatedToolResults: ToolResult[]) => void
}

export function useChatStreaming(): UseChatStreamingReturn {
  // Streaming states
  const [isStreaming, setIsStreaming] = useState(false)
  const [streamingMessage, setStreamingMessage] = useState('')
  const [aiPlanning, setAiPlanning] = useState('')
  const [currentToolExecution, setCurrentToolExecution] = useState<{
    name: string
    description: string
    args?: any
  } | null>(null)
  const [toolExecutionResults, setToolExecutionResults] = useState<Array<{
    name: string
    result: string
    success: boolean
  }>>([])

  // Refs for reliable streaming state management
  const abortControllerRef = useRef<AbortController | null>(null)
  const currentStreamingContentRef = useRef<string>('')
  // Remove streamingActivityRef - it duplicates React state

  // Control functions
  const startStreaming = useCallback(() => {
    setIsStreaming(true)
    setStreamingMessage('')
    setAiPlanning('')
    setCurrentToolExecution(null)
    setToolExecutionResults([])
    currentStreamingContentRef.current = ''

    // Create new abort controller
    abortControllerRef.current = new AbortController()
  }, [])

  const stopStreaming = useCallback(() => {
    setIsStreaming(false)
  }, [])

  const resetStreaming = useCallback(() => {
    setIsStreaming(false)
    setStreamingMessage('')
    setAiPlanning('')
    setCurrentToolExecution(null)
    setToolExecutionResults([])
    currentStreamingContentRef.current = ''
    
    // Clear abort controller
    abortControllerRef.current = null
  }, [])

  const updateStreamingContent = useCallback((content: string) => {
    setStreamingMessage(content)
    currentStreamingContentRef.current = content
  }, [])

  const updateAiPlanning = useCallback((planning: string) => {
    console.log('Setting AI planning:', planning)
    setAiPlanning(planning)
  }, [])

  const updateCurrentExecution = useCallback((execution: {
    name: string
    description: string
    args?: any
  } | null) => {
    if (execution) {
      console.log('Setting current tool execution:', execution.name, execution.description)
    }
    setCurrentToolExecution(execution)
  }, [])

  const addToolResult = useCallback((result: {
    name: string
    result: string
    success: boolean
  }) => {
    console.log('Adding tool result:', result.name, result.result)
    setCurrentToolExecution(null) // Clear current execution
    
    setToolExecutionResults(prev => {
      const newResults = [...prev, result]
      console.log('Updated tool results:', newResults)
      return newResults
    })
  }, [])

  const stopExecution = useCallback(() => {
    if (abortControllerRef.current) {
      console.log('Stopping AI execution, current streaming message:', streamingMessage)
      abortControllerRef.current.abort()
      abortControllerRef.current = null
    }
  }, [streamingMessage])

  const processStreamChunk = useCallback((parsed: any, accumulatedToolResults: ToolResult[]) => {
    if (parsed.type) {
      switch (parsed.type) {
        case 'thinking':
          // Show thinking state to user
          updateAiPlanning(`🤔 ${parsed.content || 'Thinking...'}`)
          break
          
        case 'planning':
          updateAiPlanning(parsed.content || '')
          break
          
        case 'executing':
          const execution = {
            name: parsed.tool_name || 'Unknown Tool',
            description: parsed.content || 'Executing...',
            args: parsed.tool_args
          }
          updateCurrentExecution(execution)
          break
          
        case 'tool_result':
          const toolResult = {
            name: parsed.tool_name || 'Unknown Tool',
            result: parsed.content || 'Completed',
            success: !parsed.content?.includes('Error')
          }
          
          addToolResult(toolResult)
          
          // Store for legacy compatibility
          if (parsed.tool_result) {
            accumulatedToolResults.push({
              role: 'tool',
              tool_call_id: parsed.tool_name || 'unknown',
              name: parsed.tool_name || 'unknown',
              content: JSON.stringify(parsed.tool_result)
            } as ToolResult)
          }
          break
          
        case 'content':
          // Add proper line breaks when transitioning from tool phases to content
          let newContent = currentStreamingContentRef.current || ''
          
          // If we're starting content after tools/planning, add line breaks
          if (newContent === '' && (aiPlanning || currentToolExecution || toolExecutionResults.length > 0)) {
            newContent += '\n\n'
          }
          
          newContent += (parsed.content || '')
          updateStreamingContent(newContent)
          break
          
        case 'error':
          console.error('Stream error:', parsed.content)
          updateCurrentExecution(null)
          updateAiPlanning('')
          break
          
        case 'status':
          if (parsed.status === 'completed') {
            console.log('Stream completed, final states:', {
              planning: aiPlanning,
              toolResults: toolExecutionResults,
              currentExecution: currentToolExecution
            })
          } else if (parsed.content) {
            console.log('Status update:', parsed.content)
          }
          break
      }
    }
  }, [updateAiPlanning, updateCurrentExecution, addToolResult, updateStreamingContent, aiPlanning, toolExecutionResults, currentToolExecution])

  return {
    // States
    isStreaming,
    streamingMessage,
    aiPlanning,
    currentToolExecution,
    toolExecutionResults,

    // Control functions
    startStreaming,
    stopStreaming,
    resetStreaming,
    updateStreamingContent,
    updateAiPlanning,
    updateCurrentExecution,
    addToolResult,

    // Refs
    abortControllerRef,
    currentStreamingContentRef,

    // Utilities
    stopExecution,
    processStreamChunk
  }
}