1  /**
  2   * Transforms SDK messages for streamlined output mode.
  3   *
  4   * Streamlined mode is a "distillation-resistant" output format that:
  5   * - Keeps text messages intact
  6   * - Summarizes tool calls with cumulative counts (resets when text appears)
  7   * - Omits thinking content
  8   * - Strips tool list and model info from init messages
  9   */
 10  
 11  import type { SDKAssistantMessage } from 'src/entrypoints/agentSdkTypes.js'
 12  import type { StdoutMessage } from 'src/entrypoints/sdk/controlTypes.js'
 13  import { FILE_EDIT_TOOL_NAME } from 'src/tools/FileEditTool/constants.js'
 14  import { FILE_READ_TOOL_NAME } from 'src/tools/FileReadTool/prompt.js'
 15  import { FILE_WRITE_TOOL_NAME } from 'src/tools/FileWriteTool/prompt.js'
 16  import { GLOB_TOOL_NAME } from 'src/tools/GlobTool/prompt.js'
 17  import { GREP_TOOL_NAME } from 'src/tools/GrepTool/prompt.js'
 18  import { LIST_MCP_RESOURCES_TOOL_NAME } from 'src/tools/ListMcpResourcesTool/prompt.js'
 19  import { LSP_TOOL_NAME } from 'src/tools/LSPTool/prompt.js'
 20  import { NOTEBOOK_EDIT_TOOL_NAME } from 'src/tools/NotebookEditTool/constants.js'
 21  import { TASK_STOP_TOOL_NAME } from 'src/tools/TaskStopTool/prompt.js'
 22  import { WEB_SEARCH_TOOL_NAME } from 'src/tools/WebSearchTool/prompt.js'
 23  import { extractTextContent } from 'src/utils/messages.js'
 24  import { SHELL_TOOL_NAMES } from 'src/utils/shell/shellToolUtils.js'
 25  import { capitalize } from 'src/utils/stringUtils.js'
 26  
 27  type ToolCounts = {
 28    searches: number
 29    reads: number
 30    writes: number
 31    commands: number
 32    other: number
 33  }
 34  
 35  /**
 36   * Tool categories for summarization.
 37   */
 38  const SEARCH_TOOLS = [
 39    GREP_TOOL_NAME,
 40    GLOB_TOOL_NAME,
 41    WEB_SEARCH_TOOL_NAME,
 42    LSP_TOOL_NAME,
 43  ]
 44  const READ_TOOLS = [FILE_READ_TOOL_NAME, LIST_MCP_RESOURCES_TOOL_NAME]
 45  const WRITE_TOOLS = [
 46    FILE_WRITE_TOOL_NAME,
 47    FILE_EDIT_TOOL_NAME,
 48    NOTEBOOK_EDIT_TOOL_NAME,
 49  ]
 50  const COMMAND_TOOLS = [...SHELL_TOOL_NAMES, 'Tmux', TASK_STOP_TOOL_NAME]
 51  
 52  function categorizeToolName(toolName: string): keyof ToolCounts {
 53    if (SEARCH_TOOLS.some(t => toolName.startsWith(t))) return 'searches'
 54    if (READ_TOOLS.some(t => toolName.startsWith(t))) return 'reads'
 55    if (WRITE_TOOLS.some(t => toolName.startsWith(t))) return 'writes'
 56    if (COMMAND_TOOLS.some(t => toolName.startsWith(t))) return 'commands'
 57    return 'other'
 58  }
 59  
 60  function createEmptyToolCounts(): ToolCounts {
 61    return {
 62      searches: 0,
 63      reads: 0,
 64      writes: 0,
 65      commands: 0,
 66      other: 0,
 67    }
 68  }
 69  
 70  /**
 71   * Generate a summary text for tool counts.
 72   */
 73  function getToolSummaryText(counts: ToolCounts): string | undefined {
 74    const parts: string[] = []
 75  
 76    // Use similar phrasing to collapseReadSearch.ts
 77    if (counts.searches > 0) {
 78      parts.push(
 79        `searched ${counts.searches} ${counts.searches === 1 ? 'pattern' : 'patterns'}`,
 80      )
 81    }
 82    if (counts.reads > 0) {
 83      parts.push(`read ${counts.reads} ${counts.reads === 1 ? 'file' : 'files'}`)
 84    }
 85    if (counts.writes > 0) {
 86      parts.push(
 87        `wrote ${counts.writes} ${counts.writes === 1 ? 'file' : 'files'}`,
 88      )
 89    }
 90    if (counts.commands > 0) {
 91      parts.push(
 92        `ran ${counts.commands} ${counts.commands === 1 ? 'command' : 'commands'}`,
 93      )
 94    }
 95    if (counts.other > 0) {
 96      parts.push(`${counts.other} other ${counts.other === 1 ? 'tool' : 'tools'}`)
 97    }
 98  
 99    if (parts.length === 0) {
100      return undefined
101    }
102  
103    return capitalize(parts.join(', '))
104  }
105  
106  /**
107   * Count tool uses in an assistant message and add to existing counts.
108   */
109  function accumulateToolUses(
110    message: SDKAssistantMessage,
111    counts: ToolCounts,
112  ): void {
113    const content = message.message.content
114    if (!Array.isArray(content)) {
115      return
116    }
117  
118    for (const block of content) {
119      if (block.type === 'tool_use' && 'name' in block) {
120        const category = categorizeToolName(block.name as string)
121        counts[category]++
122      }
123    }
124  }
125  
126  /**
127   * Create a stateful transformer that accumulates tool counts between text messages.
128   * Tool counts reset when a message with text content is encountered.
129   */
130  export function createStreamlinedTransformer(): (
131    message: StdoutMessage,
132  ) => StdoutMessage | null {
133    let cumulativeCounts = createEmptyToolCounts()
134  
135    return function transformToStreamlined(
136      message: StdoutMessage,
137    ): StdoutMessage | null {
138      switch (message.type) {
139        case 'assistant': {
140          const content = message.message.content
141          const text = Array.isArray(content)
142            ? extractTextContent(content, '\n').trim()
143            : ''
144  
145          // Accumulate tool counts from this message
146          accumulateToolUses(message, cumulativeCounts)
147  
148          if (text.length > 0) {
149            // Text message: emit text only, reset counts
150            cumulativeCounts = createEmptyToolCounts()
151            return {
152              type: 'streamlined_text',
153              text,
154              session_id: message.session_id,
155              uuid: message.uuid,
156            }
157          }
158  
159          // Tool-only message: emit cumulative tool summary
160          const toolSummary = getToolSummaryText(cumulativeCounts)
161          if (!toolSummary) {
162            return null
163          }
164  
165          return {
166            type: 'streamlined_tool_use_summary',
167            tool_summary: toolSummary,
168            session_id: message.session_id,
169            uuid: message.uuid,
170          }
171        }
172  
173        case 'result':
174          // Keep result messages as-is (they have structured_output, permission_denials)
175          return message
176  
177        case 'system':
178        case 'user':
179        case 'stream_event':
180        case 'tool_progress':
181        case 'auth_status':
182        case 'rate_limit_event':
183        case 'control_response':
184        case 'control_request':
185        case 'control_cancel_request':
186        case 'keep_alive':
187          return null
188  
189        default:
190          return null
191      }
192    }
193  }
194  
195  /**
196   * Check if a message should be included in streamlined output.
197   * Useful for filtering before transformation.
198   */
199  export function shouldIncludeInStreamlined(message: StdoutMessage): boolean {
200    return message.type === 'assistant' || message.type === 'result'
201  }