[codex] Summarize collapsed code searcher (#645)

Author: jahooma

cli/src/components/blocks/agent-branch-wrapper.tsx

@@ -23,7 +23,11 @@ import {
   processBlocks,
   type BlockProcessorHandlers,
 } from '../../utils/block-processor'
-import { shouldRenderAsSimpleText, isMultiPromptEditor } from '../../utils/constants'
+import { getCodeSearcherCollapsedPreview } from '../../utils/code-search-summary'
+import {
+  shouldRenderAsSimpleText,
+  isMultiPromptEditor,
+} from '../../utils/constants'
 import {
   isImplementorAgent,
   getImplementorIndex,
@@ -65,6 +69,11 @@ function getCollapsedPreview(
     }
   }
 
+  const codeSearcherPreview = getCodeSearcherCollapsedPreview(agentBlock)
+  if (codeSearcherPreview) {
+    return codeSearcherPreview
+  }
+
   // Default preview: use the displayed prompt or first line of text content.
   const displayPrompt = getAgentDisplayPrompt(agentBlock)
   if (displayPrompt) {
@@ -357,8 +366,12 @@ export const AgentBranchWrapper = memo(
             b.type === 'tool' && b.toolName === 'set_output',
         )
         // set_output wraps data in a 'data' property, so we need to access input.data
-        const outputData = (setOutputBlock?.input as { data?: Record<string, unknown> })?.data
-        const implementationId = outputData?.implementationId as string | undefined
+        const outputData = (
+          setOutputBlock?.input as { data?: Record<string, unknown> }
+        )?.data
+        const implementationId = outputData?.implementationId as
+          | string
+          | undefined
         if (implementationId) {
           const letterIndex = implementationId.charCodeAt(0) - 65
           const implementors = siblingBlocks.filter(

cli/src/components/tools/code-search.tsx

@@ -2,6 +2,7 @@ import React from 'react'
 
 import { SimpleToolCallItem } from './tool-call-item'
 import { defineToolComponent } from './types'
+import { countCodeSearchResults } from '../../utils/code-search-summary'
 
 import type { ToolRenderConfig } from './types'
 
@@ -18,30 +19,7 @@ export const CodeSearchComponent = defineToolComponent({
     const pattern = input?.pattern ?? ''
     const cwd = input?.cwd ?? ''
 
-    // Count results from output
-    let totalResults = 0
-
-    if (toolBlock.output && typeof toolBlock.output === 'string') {
-      const lines = toolBlock.output.split('\n')
-      const matchCountLine = lines.find((line) =>
-        /^Found \d+ matches?$/.test(line.trim()),
-      )
-      const parsedTotalResults = matchCountLine
-        ?.trim()
-        .match(/^Found (\d+) matches?$/)?.[1]
-
-      if (parsedTotalResults !== undefined) {
-        totalResults = Number(parsedTotalResults)
-      } else {
-        for (const line of lines) {
-          const trimmed = line.trim()
-
-          if (/^(?:Line\s+)?\d+:/.test(trimmed)) {
-            totalResults++
-          }
-        }
-      }
-    }
+    const totalResults = countCodeSearchResults(toolBlock.output)
 
     // Build single-line summary
     let summary = ''

cli/src/utils/__tests__/code-search-summary.test.ts

@@ -0,0 +1,84 @@
+import { describe, expect, test } from 'bun:test'
+
+import {
+  countCodeSearchResults,
+  getCodeSearcherCollapsedPreview,
+} from '../code-search-summary'
+
+import type { AgentContentBlock, ToolContentBlock } from '../../types/chat'
+
+const createCodeSearchToolBlock = (
+  output: string,
+  id = 'tool-1',
+): ToolContentBlock => ({
+  type: 'tool',
+  toolCallId: id,
+  toolName: 'code_search',
+  input: { pattern: 'MODEL_ID' },
+  output,
+})
+
+const createCodeSearcherBlock = (
+  options: Partial<AgentContentBlock> = {},
+): AgentContentBlock => ({
+  type: 'agent',
+  agentId: 'agent-1',
+  agentName: 'code-searcher',
+  agentType: 'code-searcher',
+  content: '',
+  status: 'complete',
+  params: {
+    searchQueries: [
+      { pattern: 'FREEBUFF_MODEL_SELECTOR_MODELS' },
+      { pattern: 'FREEBUFF_MODEL_SELECTOR_MODEL_IDS' },
+      { pattern: 'DEFAULT_FREEBUFF_MODEL_ID' },
+    ],
+  },
+  blocks: [],
+  ...options,
+})
+
+describe('code search summary helpers', () => {
+  test('counts formatted code search matches from stdout', () => {
+    expect(
+      countCodeSearchResults(`stdout: |-
+  Found 2 matches
+  ./message-block-helpers.ts:
+    Line 13: export const getAgentBaseName = (type: string): string => {
+    Line 196: getAgentBaseName(options.agentType ?? '') === 'code-searcher'`),
+    ).toBe(2)
+  })
+
+  test('summarizes collapsed code-searcher searches and results', () => {
+    const agentBlock = createCodeSearcherBlock({
+      blocks: [
+        createCodeSearchToolBlock('Found 7 matches', 'tool-1'),
+        createCodeSearchToolBlock('Found 2 matches', 'tool-2'),
+        createCodeSearchToolBlock('Found 7 matches', 'tool-3'),
+      ],
+    })
+
+    expect(getCodeSearcherCollapsedPreview(agentBlock)).toBe(
+      '3 searches · 16 results',
+    )
+  })
+
+  test('shows search count before tool outputs arrive', () => {
+    expect(getCodeSearcherCollapsedPreview(createCodeSearcherBlock())).toBe(
+      '3 searches',
+    )
+  })
+
+  test('handles singular labels', () => {
+    const agentBlock = createCodeSearcherBlock({
+      params: {
+        searchQueries: [{ pattern: 'DEFAULT_FREEBUFF_MODEL_ID' }],
+      },
+      blocks: [createCodeSearchToolBlock('Found 1 match')],
+    })
+
+    expect(getCodeSearcherCollapsedPreview(agentBlock)).toBe(
+      '1 search · 1 result',
+    )
+  })
+})

cli/src/utils/code-search-summary.ts

@@ -0,0 +1,70 @@
+import { getAgentBaseName } from './message-block-helpers'
+
+import type {
+  AgentContentBlock,
+  ContentBlock,
+  ToolContentBlock,
+} from '../types/chat'
+
+export function countCodeSearchResults(output?: string): number {
+  if (!output) {
+    return 0
+  }
+
+  const lines = output.split('\n')
+  const matchCountLine = lines.find((line) =>
+    /^Found \d+ match(?:es)?$/.test(line.trim()),
+  )
+  const parsedTotalResults = matchCountLine
+    ?.trim()
+    .match(/^Found (\d+) match(?:es)?$/)?.[1]
+
+  if (parsedTotalResults !== undefined) {
+    return Number(parsedTotalResults)
+  }
+
+  return lines.reduce((total, line) => {
+    const trimmed = line.trim()
+    return /^(?:Line\s+)?\d+:/.test(trimmed) ? total + 1 : total
+  }, 0)
+}
+
+const pluralize = (count: number, singular: string, plural = `${singular}s`) =>
+  `${count} ${count === 1 ? singular : plural}`
+
+const isCodeSearchToolBlock = (
+  block: ContentBlock,
+): block is ToolContentBlock =>
+  block.type === 'tool' && block.toolName === 'code_search'
+
+export function getCodeSearcherCollapsedPreview(
+  agentBlock: AgentContentBlock,
+): string | undefined {
+  if (getAgentBaseName(agentBlock.agentType) !== 'code-searcher') {
+    return undefined
+  }
+
+  const toolBlocks = (agentBlock.blocks ?? []).filter(isCodeSearchToolBlock)
+  const searchQueries = Array.isArray(agentBlock.params?.searchQueries)
+    ? agentBlock.params.searchQueries
+    : []
+  const searchCount = searchQueries.length || toolBlocks.length
+
+  if (searchCount === 0) {
+    return undefined
+  }
+
+  const completedToolBlocks = toolBlocks.filter((block) => block.output)
+  const searchLabel = pluralize(searchCount, 'search', 'searches')
+
+  if (completedToolBlocks.length === 0) {
+    return searchLabel
+  }
+
+  const totalResults = completedToolBlocks.reduce(
+    (total, block) => total + countCodeSearchResults(block.output),
+    0,
+  )
+
+  return `${searchLabel} · ${pluralize(totalResults, 'result')}`
+}