1  import type {
 2    AssistantMessage,
 3    AttachmentMessage,
 4    SystemMessage,
 5    UserMessage,
 6  } from 'src/types/message.js'
 7  
 8  /**
 9   * Tags user messages with a sourceToolUseID so they stay transient until the tool resolves.
10   * This prevents the "is running" message from being duplicated in the UI.
11   */
12  export function tagMessagesWithToolUseID(
13    messages: (UserMessage | AttachmentMessage | SystemMessage)[],
14    toolUseID: string | undefined,
15  ): (UserMessage | AttachmentMessage | SystemMessage)[] {
16    if (!toolUseID) {
17      return messages
18    }
19    return messages.map(m => {
20      if (m.type === 'user') {
21        return { ...m, sourceToolUseID: toolUseID }
22      }
23      return m
24    })
25  }
26  
27  /**
28   * Extracts the tool use ID from a parent message for a given tool name.
29   */
30  export function getToolUseIDFromParentMessage(
31    parentMessage: AssistantMessage,
32    toolName: string,
33  ): string | undefined {
34    const toolUseBlock = parentMessage.message.content.find(
35      block => block.type === 'tool_use' && block.name === toolName,
36    )
37    return toolUseBlock && toolUseBlock.type === 'tool_use'
38      ? toolUseBlock.id
39      : undefined
40  }