1  import { Box } from 'ink'
  2  import * as React from 'react'
  3  import type { AssistantMessage, Message, UserMessage } from '../query.js'
  4  import type {
  5    ContentBlock,
  6    DocumentBlockParam,
  7    ImageBlockParam,
  8    TextBlockParam,
  9    ThinkingBlockParam,
 10    ToolResultBlockParam,
 11    ToolUseBlockParam,
 12  } from '@anthropic-ai/sdk/resources/index.mjs'
 13  import { Tool } from '../Tool.js'
 14  import { logError } from '../utils/log.js'
 15  import { UserToolResultMessage } from './messages/UserToolResultMessage/UserToolResultMessage.js'
 16  import { AssistantToolUseMessage } from './messages/AssistantToolUseMessage.js'
 17  import { AssistantTextMessage } from './messages/AssistantTextMessage.js'
 18  import { UserTextMessage } from './messages/UserTextMessage.js'
 19  import { NormalizedMessage } from '../utils/messages.js'
 20  import { AssistantThinkingMessage } from './messages/AssistantThinkingMessage.js'
 21  import { AssistantRedactedThinkingMessage } from './messages/AssistantRedactedThinkingMessage.js'
 22  import { useTerminalSize } from '../hooks/useTerminalSize.js'
 23  
 24  type Props = {
 25    message: UserMessage | AssistantMessage
 26    messages: NormalizedMessage[]
 27    // TODO: Find a way to remove this, and leave spacing to the consumer
 28    addMargin: boolean
 29    tools: Tool[]
 30    verbose: boolean
 31    debug: boolean
 32    erroredToolUseIDs: Set<string>
 33    inProgressToolUseIDs: Set<string>
 34    unresolvedToolUseIDs: Set<string>
 35    shouldAnimate: boolean
 36    shouldShowDot: boolean
 37    width?: number | string
 38  }
 39  
 40  export function Message({
 41    message,
 42    messages,
 43    addMargin,
 44    tools,
 45    verbose,
 46    debug,
 47    erroredToolUseIDs,
 48    inProgressToolUseIDs,
 49    unresolvedToolUseIDs,
 50    shouldAnimate,
 51    shouldShowDot,
 52    width,
 53  }: Props): React.ReactNode {
 54    // Assistant message
 55    if (message.type === 'assistant') {
 56      return (
 57        <Box flexDirection="column" width="100%">
 58          {message.message.content.map((_, index) => (
 59            <AssistantMessage
 60              key={index}
 61              param={_}
 62              costUSD={message.costUSD}
 63              durationMs={message.durationMs}
 64              addMargin={addMargin}
 65              tools={tools}
 66              debug={debug}
 67              options={{ verbose }}
 68              erroredToolUseIDs={erroredToolUseIDs}
 69              inProgressToolUseIDs={inProgressToolUseIDs}
 70              unresolvedToolUseIDs={unresolvedToolUseIDs}
 71              shouldAnimate={shouldAnimate}
 72              shouldShowDot={shouldShowDot}
 73              width={width}
 74            />
 75          ))}
 76        </Box>
 77      )
 78    }
 79  
 80    // User message
 81    // TODO: normalize upstream
 82    const content =
 83      typeof message.message.content === 'string'
 84        ? [{ type: 'text', text: message.message.content } as TextBlockParam]
 85        : message.message.content
 86    return (
 87      <Box flexDirection="column" width="100%">
 88        {content.map((_, index) => (
 89          <UserMessage
 90            key={index}
 91            message={message}
 92            messages={messages}
 93            addMargin={addMargin}
 94            tools={tools}
 95            param={_ as TextBlockParam}
 96            options={{ verbose }}
 97          />
 98        ))}
 99      </Box>
100    )
101  }
102  
103  function UserMessage({
104    message,
105    messages,
106    addMargin,
107    tools,
108    param,
109    options: { verbose },
110  }: {
111    message: UserMessage
112    messages: Message[]
113    addMargin: boolean
114    tools: Tool[]
115    param:
116      | TextBlockParam
117      | DocumentBlockParam
118      | ImageBlockParam
119      | ToolUseBlockParam
120      | ToolResultBlockParam
121    options: {
122      verbose: boolean
123    }
124  }): React.ReactNode {
125    const { columns } = useTerminalSize()
126    switch (param.type) {
127      case 'text':
128        return <UserTextMessage addMargin={addMargin} param={param} />
129      case 'tool_result':
130        return (
131          <UserToolResultMessage
132            param={param}
133            message={message}
134            messages={messages}
135            tools={tools}
136            verbose={verbose}
137            width={columns - 5}
138          />
139        )
140    }
141  }
142  
143  function AssistantMessage({
144    param,
145    costUSD,
146    durationMs,
147    addMargin,
148    tools,
149    debug,
150    options: { verbose },
151    erroredToolUseIDs,
152    inProgressToolUseIDs,
153    unresolvedToolUseIDs,
154    shouldAnimate,
155    shouldShowDot,
156    width,
157  }: {
158    param:
159      | ContentBlock
160      | TextBlockParam
161      | ImageBlockParam
162      | ThinkingBlockParam
163      | ToolUseBlockParam
164      | ToolResultBlockParam
165    costUSD: number
166    durationMs: number
167    addMargin: boolean
168    tools: Tool[]
169    debug: boolean
170    options: {
171      verbose: boolean
172    }
173    erroredToolUseIDs: Set<string>
174    inProgressToolUseIDs: Set<string>
175    unresolvedToolUseIDs: Set<string>
176    shouldAnimate: boolean
177    shouldShowDot: boolean
178    width?: number | string
179  }): React.ReactNode {
180    switch (param.type) {
181      case 'tool_use':
182        return (
183          <AssistantToolUseMessage
184            param={param}
185            costUSD={costUSD}
186            durationMs={durationMs}
187            addMargin={addMargin}
188            tools={tools}
189            debug={debug}
190            verbose={verbose}
191            erroredToolUseIDs={erroredToolUseIDs}
192            inProgressToolUseIDs={inProgressToolUseIDs}
193            unresolvedToolUseIDs={unresolvedToolUseIDs}
194            shouldAnimate={shouldAnimate}
195            shouldShowDot={shouldShowDot}
196          />
197        )
198      case 'text':
199        return (
200          <AssistantTextMessage
201            param={param}
202            costUSD={costUSD}
203            durationMs={durationMs}
204            debug={debug}
205            addMargin={addMargin}
206            shouldShowDot={shouldShowDot}
207            verbose={verbose}
208            width={width}
209          />
210        )
211      case 'redacted_thinking':
212        return <AssistantRedactedThinkingMessage addMargin={addMargin} />
213      case 'thinking':
214        return <AssistantThinkingMessage addMargin={addMargin} param={param} />
215      default:
216        logError(`Unable to render message type: ${param.type}`)
217        return null
218    }
219  }