1  import { Box, Text, useInput } from 'ink'
  2  import * as React from 'react'
  3  import { useMemo, useState, useEffect } from 'react'
  4  import figures from 'figures'
  5  import { getTheme } from '../utils/theme.js'
  6  import { Message as MessageComponent } from './Message.js'
  7  import { randomUUID } from 'crypto'
  8  import { type Tool } from '../Tool.js'
  9  import {
 10    createUserMessage,
 11    isEmptyMessageText,
 12    isNotEmptyMessage,
 13    normalizeMessages,
 14  } from '../utils/messages.js'
 15  import { logEvent } from '../services/statsig.js'
 16  import type { AssistantMessage, UserMessage } from '../query.js'
 17  import { useExitOnCtrlCD } from '../hooks/useExitOnCtrlCD.js'
 18  
 19  type Props = {
 20    erroredToolUseIDs: Set<string>
 21    messages: (UserMessage | AssistantMessage)[]
 22    onSelect: (message: UserMessage) => void
 23    onEscape: () => void
 24    tools: Tool[]
 25    unresolvedToolUseIDs: Set<string>
 26  }
 27  
 28  const MAX_VISIBLE_MESSAGES = 7
 29  
 30  export function MessageSelector({
 31    erroredToolUseIDs,
 32    messages,
 33    onSelect,
 34    onEscape,
 35    tools,
 36    unresolvedToolUseIDs,
 37  }: Props): React.ReactNode {
 38    const currentUUID = useMemo(randomUUID, [])
 39  
 40    // Log when selector is opened
 41    useEffect(() => {
 42      logEvent('tengu_message_selector_opened', {})
 43    }, [])
 44  
 45    function handleSelect(message: UserMessage) {
 46      const indexFromEnd = messages.length - 1 - messages.indexOf(message)
 47      logEvent('tengu_message_selector_selected', {
 48        index_from_end: indexFromEnd.toString(),
 49        message_type: message.type,
 50        is_current_prompt: (message.uuid === currentUUID).toString(),
 51      })
 52      onSelect(message)
 53    }
 54  
 55    function handleEscape() {
 56      logEvent('tengu_message_selector_cancelled', {})
 57      onEscape()
 58    }
 59  
 60    // Add current prompt as a virtual message
 61    const allItems = useMemo(
 62      () => [
 63        // Filter out tool results
 64        ...messages
 65          .filter(
 66            _ =>
 67              !(
 68                _.type === 'user' &&
 69                Array.isArray(_.message.content) &&
 70                _.message.content[0]?.type === 'tool_result'
 71              ),
 72          )
 73          // Filter out assistant messages, until we have a way to kick off the tool use loop from REPL
 74          .filter(_ => _.type !== 'assistant'),
 75        { ...createUserMessage(''), uuid: currentUUID } as UserMessage,
 76      ],
 77      [messages, currentUUID],
 78    )
 79    const [selectedIndex, setSelectedIndex] = useState(allItems.length - 1)
 80  
 81    const exitState = useExitOnCtrlCD(() => process.exit(0))
 82  
 83    useInput((input, key) => {
 84      if (key.tab || key.escape) {
 85        handleEscape()
 86        return
 87      }
 88      if (key.return) {
 89        handleSelect(allItems[selectedIndex]!)
 90        return
 91      }
 92      if (key.upArrow) {
 93        if (key.ctrl || key.shift || key.meta) {
 94          // Jump to top with any modifier key
 95          setSelectedIndex(0)
 96        } else {
 97          setSelectedIndex(prev => Math.max(0, prev - 1))
 98        }
 99      }
100      if (key.downArrow) {
101        if (key.ctrl || key.shift || key.meta) {
102          // Jump to bottom with any modifier key
103          setSelectedIndex(allItems.length - 1)
104        } else {
105          setSelectedIndex(prev => Math.min(allItems.length - 1, prev + 1))
106        }
107      }
108  
109      // Handle number keys (1-9)
110      const num = Number(input)
111      if (!isNaN(num) && num >= 1 && num <= Math.min(9, allItems.length)) {
112        if (!allItems[num - 1]) {
113          return
114        }
115        handleSelect(allItems[num - 1]!)
116      }
117    })
118  
119    const firstVisibleIndex = Math.max(
120      0,
121      Math.min(
122        selectedIndex - Math.floor(MAX_VISIBLE_MESSAGES / 2),
123        allItems.length - MAX_VISIBLE_MESSAGES,
124      ),
125    )
126  
127    const normalizedMessages = useMemo(
128      () => normalizeMessages(messages).filter(isNotEmptyMessage),
129      [messages],
130    )
131  
132    return (
133      <>
134        <Box
135          flexDirection="column"
136          borderStyle="round"
137          borderColor={getTheme().secondaryBorder}
138          height={4 + Math.min(MAX_VISIBLE_MESSAGES, allItems.length) * 2}
139          paddingX={1}
140          marginTop={1}
141        >
142          <Box flexDirection="column" minHeight={2} marginBottom={1}>
143            <Text bold>Jump to a previous message</Text>
144            <Text dimColor>This will fork the conversation</Text>
145          </Box>
146          {allItems
147            .slice(firstVisibleIndex, firstVisibleIndex + MAX_VISIBLE_MESSAGES)
148            .map((msg, index) => {
149              const actualIndex = firstVisibleIndex + index
150              const isSelected = actualIndex === selectedIndex
151              const isCurrent = msg.uuid === currentUUID
152  
153              return (
154                <Box key={msg.uuid} flexDirection="row" height={2} minHeight={2}>
155                  <Box width={7}>
156                    {isSelected ? (
157                      <Text color="blue" bold>
158                        {figures.pointer} {firstVisibleIndex + index + 1}{' '}
159                      </Text>
160                    ) : (
161                      <Text>
162                        {'  '}
163                        {firstVisibleIndex + index + 1}{' '}
164                      </Text>
165                    )}
166                  </Box>
167                  <Box height={1} overflow="hidden" width={100}>
168                    {isCurrent ? (
169                      <Box width="100%">
170                        <Text dimColor italic>
171                          {'(current)'}
172                        </Text>
173                      </Box>
174                    ) : Array.isArray(msg.message.content) &&
175                      msg.message.content[0]?.type === 'text' &&
176                      isEmptyMessageText(msg.message.content[0].text) ? (
177                      <Text dimColor italic>
178                        (empty message)
179                      </Text>
180                    ) : (
181                      <MessageComponent
182                        message={msg}
183                        messages={normalizedMessages}
184                        addMargin={false}
185                        tools={tools}
186                        verbose={false}
187                        debug={false}
188                        erroredToolUseIDs={erroredToolUseIDs}
189                        inProgressToolUseIDs={new Set()}
190                        unresolvedToolUseIDs={unresolvedToolUseIDs}
191                        shouldAnimate={false}
192                        shouldShowDot={false}
193                      />
194                    )}
195                  </Box>
196                </Box>
197              )
198            })}
199        </Box>
200        <Box marginLeft={3}>
201          <Text dimColor>
202            {exitState.pending ? (
203              <>Press {exitState.keyName} again to exit</>
204            ) : (
205              <>↑/↓ to select · Enter to confirm · Tab/Esc to cancel</>
206            )}
207          </Text>
208        </Box>
209      </>
210    )
211  }