deepspace 0.6.2 → 0.8.0

package/features/ai-chat/src/ChatPanel.stream.ts

@@ -0,0 +1,397 @@
+import {
+  useCallback,
+  useEffect,
+  useRef,
+  useState,
+  type Dispatch,
+  type SetStateAction,
+} from 'react'
+import {
+  decodeAiStreamChunk,
+  getAuthToken,
+  parseSseLine,
+  type AiStreamAction,
+} from 'deepspace'
+
+export type ToolInvocation = {
+  toolName: string
+  state: 'call' | 'result'
+  args?: unknown
+  result?: unknown
+}
+
+export type ToolInvocationPart = {
+  type: 'tool-invocation'
+  toolInvocation: ToolInvocation
+  toolCallId: string
+}
+
+export type MessagePart =
+  | { type: 'text'; text: string }
+  | ToolInvocationPart
+
+export type InFlightMessage = {
+  id: string
+  role: 'user' | 'assistant'
+  content: string
+  parts: MessagePart[]
+  /** The persisted assistant record id returned in `X-Asst-Id`. */
+  serverId?: string
+  /** Prevents an old chat's overlay from leaking into the active chat. */
+  forChatId: string
+}
+
+type UseStreamingChatOptions = {
+  chatId: string | null
+  modelId?: string
+  onChatCreated?: (id: string) => void
+}
+
+/**
+ * Applies one decoded stream action to the in-flight overlay.
+ *
+ * This is deliberately pure: the SDK owns wire decoding, this reducer owns
+ * overlay semantics, and the hook below owns transport and React lifecycle.
+ */
+export function reduceStreamMessages(
+  messages: InFlightMessage[],
+  action: AiStreamAction,
+  assistantId: string,
+): InFlightMessage[] {
+  switch (action.type) {
+    case 'append-text':
+      return messages.map((message) => {
+        if (message.id !== assistantId) return message
+        const last = message.parts[message.parts.length - 1]
+        const parts: MessagePart[] = last?.type === 'text'
+          ? [
+              ...message.parts.slice(0, -1),
+              { type: 'text', text: last.text + action.delta },
+            ]
+          : [...message.parts, { type: 'text', text: action.delta }]
+        return {
+          ...message,
+          content: message.content + action.delta,
+          parts,
+        }
+      })
+
+    case 'upsert-tool-call':
+      return upsertToolInvocation(messages, assistantId, action.toolCallId, {
+        toolName: action.toolName,
+        state: 'call',
+        args: action.input,
+      })
+
+    case 'finalize-tool-call':
+      return finalizeToolInvocation(
+        messages,
+        assistantId,
+        action.toolCallId,
+        action.result,
+      )
+
+    case 'fail-tool-input':
+      return upsertToolInvocation(messages, assistantId, action.toolCallId, {
+        toolName: action.toolName,
+        state: 'result',
+        args: action.input,
+        result: { success: false, error: action.errorText },
+      })
+
+    case 'fail-tool-output':
+      return finalizeToolInvocation(messages, assistantId, action.toolCallId, {
+        success: false,
+        error: action.errorText,
+      })
+
+    case 'stream-error':
+    case 'abort':
+      // Preserve partial text or a tool row the user has already seen, but do
+      // not leave an empty failed/stopped assistant occupying layout space.
+      return messages.filter(
+        (message) => message.id !== assistantId || message.parts.length > 0,
+      )
+
+    default: {
+      const exhaustive: never = action
+      return exhaustive
+    }
+  }
+}
+
+function upsertToolInvocation(
+  messages: InFlightMessage[],
+  assistantId: string,
+  toolCallId: string,
+  invocation: ToolInvocation,
+): InFlightMessage[] {
+  const part: ToolInvocationPart = {
+    type: 'tool-invocation',
+    toolInvocation: invocation,
+    toolCallId,
+  }
+
+  return messages.map((message) => {
+    if (message.id !== assistantId) return message
+    const index = message.parts.findIndex(
+      (candidate) =>
+        candidate.type === 'tool-invocation' && candidate.toolCallId === toolCallId,
+    )
+    if (index < 0) return { ...message, parts: [...message.parts, part] }
+
+    const parts = message.parts.slice()
+    parts[index] = part
+    return { ...message, parts }
+  })
+}
+
+function finalizeToolInvocation(
+  messages: InFlightMessage[],
+  assistantId: string,
+  toolCallId: string,
+  result: unknown,
+): InFlightMessage[] {
+  return messages.map((message) => {
+    if (message.id !== assistantId) return message
+    const parts = message.parts.map((part): MessagePart => {
+      if (part.type !== 'tool-invocation' || part.toolCallId !== toolCallId) {
+        return part
+      }
+      return {
+        ...part,
+        toolInvocation: {
+          ...part.toolInvocation,
+          state: 'result',
+          result,
+        },
+      }
+    })
+    return { ...message, parts }
+  })
+}
+
+export function useStreamingChat({
+  chatId,
+  modelId,
+  onChatCreated,
+}: UseStreamingChatOptions) {
+  const [inFlight, setInFlight] = useState<InFlightMessage[]>([])
+  const [isLoading, setIsLoading] = useState(false)
+  const [error, setError] = useState<Error | null>(null)
+  const isLoadingRef = useRef(false)
+  const abortRef = useRef<AbortController | null>(null)
+  const assistantIdRef = useRef<string | null>(null)
+  const lastSendRef = useRef<string | null>(null)
+  const previousChatIdRef = useRef<string | null>(chatId)
+
+  useEffect(() => {
+    const previousChatId = previousChatIdRef.current
+    previousChatIdRef.current = chatId
+    // `null -> id` is the auto-create promotion during the first send. Every
+    // transition away from a real id is a genuine switch and aborts its turn.
+    if (previousChatId !== null && previousChatId !== chatId) {
+      abortRef.current?.abort()
+      setInFlight([])
+    }
+  }, [chatId])
+
+  useEffect(() => () => abortRef.current?.abort(), [])
+
+  const send = useCallback(
+    async (content: string) => {
+      // Claim the slot synchronously so two first sends cannot create two chats.
+      if (isLoadingRef.current) return
+      isLoadingRef.current = true
+      setIsLoading(true)
+      setError(null)
+      lastSendRef.current = content
+
+      const controller = new AbortController()
+      abortRef.current = controller
+      let assistantId: string | null = null
+      try {
+        let activeChatId = chatId
+        if (!activeChatId) {
+          activeChatId = await createChat(controller.signal)
+          if (controller.signal.aborted) return
+          onChatCreated?.(activeChatId)
+        }
+
+        const localTimestamp = Date.now()
+        const userMessageId =
+          `usr-${localTimestamp}-${Math.random().toString(36).slice(2, 8)}`
+        assistantId = `asst-pending-${localTimestamp}`
+        assistantIdRef.current = assistantId
+        setInFlight([
+          {
+            id: userMessageId,
+            role: 'user',
+            content,
+            parts: [],
+            forChatId: activeChatId,
+          },
+          {
+            id: assistantId,
+            role: 'assistant',
+            content: '',
+            parts: [],
+            forChatId: activeChatId,
+          },
+        ])
+
+        const token = await getAuthToken()
+        const response = await fetch('/api/ai/chat', {
+          method: 'POST',
+          headers: authHeaders(token),
+          body: JSON.stringify({
+            chatId: activeChatId,
+            userMessageId,
+            content,
+            modelId,
+          }),
+          signal: controller.signal,
+        })
+        if (!response.ok || !response.body) {
+          const detail = response.body
+            ? await response.text().catch(() => '')
+            : ''
+          throw new Error(detail || `Request failed: ${response.status}`)
+        }
+
+        const serverId = response.headers.get('X-Asst-Id')
+        if (serverId) {
+          setInFlight((messages) =>
+            messages.map((message) =>
+              message.id === assistantId ? { ...message, serverId } : message,
+            ),
+          )
+        }
+
+        await consumeStream(response.body, assistantId, setInFlight, setError)
+      } catch (cause) {
+        const streamError = toError(cause)
+        // A local abort does not necessarily deliver an SSE `abort` event.
+        // Apply the same overlay rule immediately for stop and transport errors.
+        if (assistantId) {
+          const failedAssistantId = assistantId
+          setInFlight((messages) =>
+            reduceStreamMessages(messages, { type: 'abort' }, failedAssistantId),
+          )
+        }
+        if (streamError.name !== 'AbortError') {
+          console.error('[chat] STREAM error', {
+            name: streamError.name,
+            message: streamError.message,
+          })
+          setError(streamError)
+        }
+      } finally {
+        if (assistantId && assistantIdRef.current === assistantId) {
+          assistantIdRef.current = null
+        }
+        if (abortRef.current === controller) abortRef.current = null
+        isLoadingRef.current = false
+        setIsLoading(false)
+      }
+    },
+    [chatId, modelId, onChatCreated],
+  )
+
+  const stop = useCallback(() => {
+    abortRef.current?.abort()
+    const assistantId = assistantIdRef.current
+    if (assistantId) {
+      setInFlight((messages) =>
+        reduceStreamMessages(messages, { type: 'abort' }, assistantId),
+      )
+    }
+  }, [])
+
+  const retry = useCallback(() => {
+    if (lastSendRef.current) void send(lastSendRef.current)
+  }, [send])
+
+  return { send, stop, retry, isLoading, error, inFlight }
+}
+
+async function createChat(signal: AbortSignal): Promise<string> {
+  const token = await getAuthToken()
+  const response = await fetch('/api/ai/chats', {
+    method: 'POST',
+    headers: authHeaders(token),
+    body: JSON.stringify({}),
+    signal,
+  })
+  if (!response.ok) throw new Error(`Failed to create chat: ${response.status}`)
+  const data = (await response.json()) as { chat?: { recordId?: unknown } }
+  if (typeof data.chat?.recordId !== 'string') {
+    throw new Error('Failed to create chat: response is missing chat.recordId')
+  }
+  return data.chat.recordId
+}
+
+function authHeaders(token: string | null): Record<string, string> {
+  return {
+    'Content-Type': 'application/json',
+    ...(token ? { Authorization: `Bearer ${token}` } : {}),
+  }
+}
+
+async function consumeStream(
+  body: ReadableStream<Uint8Array>,
+  assistantId: string,
+  setInFlight: Dispatch<SetStateAction<InFlightMessage[]>>,
+  setError: Dispatch<SetStateAction<Error | null>>,
+) {
+  const reader = body.getReader()
+  const decoder = new TextDecoder()
+  let buffer = ''
+
+  while (true) {
+    const { done, value } = await reader.read()
+    if (done) break
+    buffer += decoder.decode(value, { stream: true })
+    const lines = buffer.split('\n')
+    buffer = lines.pop() ?? ''
+    for (const line of lines) {
+      applySseLine(line, assistantId, setInFlight, setError)
+    }
+  }
+
+  buffer += decoder.decode()
+  applySseLine(buffer, assistantId, setInFlight, setError)
+}
+
+function applySseLine(
+  line: string,
+  assistantId: string,
+  setInFlight: Dispatch<SetStateAction<InFlightMessage[]>>,
+  setError: Dispatch<SetStateAction<Error | null>>,
+) {
+  const chunk = parseSseLine(line)
+  if (!chunk) return
+  const action = decodeAiStreamChunk(chunk)
+  if (!action) return
+
+  if (action.type === 'stream-error') {
+    console.error('[chat] STREAM error', action.errorText)
+    setError(new Error(action.errorText))
+  } else if (action.type === 'fail-tool-input') {
+    console.error('[chat] STREAM tool-input-error', {
+      toolCallId: action.toolCallId,
+      errorText: action.errorText,
+    })
+  } else if (action.type === 'fail-tool-output') {
+    console.error('[chat] STREAM tool-output-error', {
+      toolCallId: action.toolCallId,
+      errorText: action.errorText,
+    })
+  }
+
+  setInFlight((messages) => reduceStreamMessages(messages, action, assistantId))
+}
+
+function toError(cause: unknown): Error {
+  return cause instanceof Error ? cause : new Error(String(cause))
+}