Function createMessageUpdater

cli/src/utils/message-updater.ts:31–121

(
  aiMessageId: string,
  setMessages: SetMessagesFn,
)

export const DEFAULT_FLUSH_INTERVAL_MS = 100

export const createMessageUpdater = (
  aiMessageId: string,
  setMessages: SetMessagesFn,
): MessageUpdater => {
  const updateAiMessage = (updater: (msg: ChatMessage) => ChatMessage) => {
    setMessages((prev) =>
      prev.map((msg) => (msg.id === aiMessageId ? updater(msg) : msg)),
    )
  }

  const updateAiMessageBlocks = (
    blockUpdater: (blocks: ContentBlock[]) => ContentBlock[],
  ) => {
    updateAiMessage((msg) => ({
      ...msg,
      blocks: blockUpdater(msg.blocks ?? []),
    }))
  }

  const addBlock = (block: ContentBlock) => {
    updateAiMessage((msg) => ({
      ...msg,
      blocks: [...(msg.blocks ?? []), block],
    }))
  }

  const markComplete = (metadata?: Partial<ChatMessage>) => {
    updateAiMessage((msg) => {
      const { metadata: messageMetadata, ...rest } = metadata ?? {}
      
      // Mark native reasoning blocks as complete by setting thinkingOpen = false
      // This ensures thinking blocks auto-collapse when the message finishes
      // Check for thinkingOpen !== false to handle both true (native) and undefined (legacy)
      const updatedBlocks = msg.blocks?.map((block) => {
        if (
          block.type === 'text' &&
          (block as TextContentBlock).textType === 'reasoning' &&
          (block as TextContentBlock).thinkingOpen !== false
        ) {
          return { ...block, thinkingOpen: false } as ContentBlock
        }
        return block
      })
      
      const nextMessage: ChatMessage = {
        ...msg,
        isComplete: true,
        ...(updatedBlocks && { blocks: updatedBlocks }),
        ...rest,
      }

      if (messageMetadata) {
        nextMessage.metadata = {
          ...(msg.metadata ?? {}),
          ...messageMetadata,
        }
      }