/ hooks / unifiedSuggestions.ts
unifiedSuggestions.ts
  1  import Fuse from 'fuse.js'
  2  import { basename } from 'path'
  3  import type { SuggestionItem } from 'src/components/PromptInput/PromptInputFooterSuggestions.js'
  4  import { generateFileSuggestions } from 'src/hooks/fileSuggestions.js'
  5  import type { ServerResource } from 'src/services/mcp/types.js'
  6  import { getAgentColor } from 'src/tools/AgentTool/agentColorManager.js'
  7  import type { AgentDefinition } from 'src/tools/AgentTool/loadAgentsDir.js'
  8  import { truncateToWidth } from 'src/utils/format.js'
  9  import { logError } from 'src/utils/log.js'
 10  import type { Theme } from 'src/utils/theme.js'
 11  
 12  type FileSuggestionSource = {
 13    type: 'file'
 14    displayText: string
 15    description?: string
 16    path: string
 17    filename: string
 18    score?: number
 19  }
 20  
 21  type McpResourceSuggestionSource = {
 22    type: 'mcp_resource'
 23    displayText: string
 24    description: string
 25    server: string
 26    uri: string
 27    name: string
 28  }
 29  
 30  type AgentSuggestionSource = {
 31    type: 'agent'
 32    displayText: string
 33    description: string
 34    agentType: string
 35    color?: keyof Theme
 36  }
 37  
 38  type SuggestionSource =
 39    | FileSuggestionSource
 40    | McpResourceSuggestionSource
 41    | AgentSuggestionSource
 42  
 43  /**
 44   * Creates a unified suggestion item from a source
 45   */
 46  function createSuggestionFromSource(source: SuggestionSource): SuggestionItem {
 47    switch (source.type) {
 48      case 'file':
 49        return {
 50          id: `file-${source.path}`,
 51          displayText: source.displayText,
 52          description: source.description,
 53        }
 54      case 'mcp_resource':
 55        return {
 56          id: `mcp-resource-${source.server}__${source.uri}`,
 57          displayText: source.displayText,
 58          description: source.description,
 59        }
 60      case 'agent':
 61        return {
 62          id: `agent-${source.agentType}`,
 63          displayText: source.displayText,
 64          description: source.description,
 65          color: source.color,
 66        }
 67    }
 68  }
 69  
 70  const MAX_UNIFIED_SUGGESTIONS = 15
 71  const DESCRIPTION_MAX_LENGTH = 60
 72  
 73  function truncateDescription(description: string): string {
 74    return truncateToWidth(description, DESCRIPTION_MAX_LENGTH)
 75  }
 76  
 77  function generateAgentSuggestions(
 78    agents: AgentDefinition[],
 79    query: string,
 80    showOnEmpty = false,
 81  ): AgentSuggestionSource[] {
 82    if (!query && !showOnEmpty) {
 83      return []
 84    }
 85  
 86    try {
 87      const agentSources: AgentSuggestionSource[] = agents.map(agent => ({
 88        type: 'agent' as const,
 89        displayText: `${agent.agentType} (agent)`,
 90        description: truncateDescription(agent.whenToUse),
 91        agentType: agent.agentType,
 92        color: getAgentColor(agent.agentType),
 93      }))
 94  
 95      if (!query) {
 96        return agentSources
 97      }
 98  
 99      const queryLower = query.toLowerCase()
100      return agentSources.filter(
101        agent =>
102          agent.agentType.toLowerCase().includes(queryLower) ||
103          agent.displayText.toLowerCase().includes(queryLower),
104      )
105    } catch (error) {
106      logError(error as Error)
107      return []
108    }
109  }
110  
111  export async function generateUnifiedSuggestions(
112    query: string,
113    mcpResources: Record<string, ServerResource[]>,
114    agents: AgentDefinition[],
115    showOnEmpty = false,
116  ): Promise<SuggestionItem[]> {
117    if (!query && !showOnEmpty) {
118      return []
119    }
120  
121    const [fileSuggestions, agentSources] = await Promise.all([
122      generateFileSuggestions(query, showOnEmpty),
123      Promise.resolve(generateAgentSuggestions(agents, query, showOnEmpty)),
124    ])
125  
126    const fileSources: FileSuggestionSource[] = fileSuggestions.map(
127      suggestion => ({
128        type: 'file' as const,
129        displayText: suggestion.displayText,
130        description: suggestion.description,
131        path: suggestion.displayText, // Use displayText as path for files
132        filename: basename(suggestion.displayText),
133        score: (suggestion.metadata as { score?: number } | undefined)?.score,
134      }),
135    )
136  
137    const mcpSources: McpResourceSuggestionSource[] = Object.values(mcpResources)
138      .flat()
139      .map(resource => ({
140        type: 'mcp_resource' as const,
141        displayText: `${resource.server}:${resource.uri}`,
142        description: truncateDescription(
143          resource.description || resource.name || resource.uri,
144        ),
145        server: resource.server,
146        uri: resource.uri,
147        name: resource.name || resource.uri,
148      }))
149  
150    if (!query) {
151      const allSources = [...fileSources, ...mcpSources, ...agentSources]
152      return allSources
153        .slice(0, MAX_UNIFIED_SUGGESTIONS)
154        .map(createSuggestionFromSource)
155    }
156  
157    const nonFileSources: SuggestionSource[] = [...mcpSources, ...agentSources]
158  
159    // Score non-file sources with Fuse.js
160    // File sources are already scored by Rust/nucleo
161    type ScoredSource = { source: SuggestionSource; score: number }
162    const scoredResults: ScoredSource[] = []
163  
164    // Add file sources with their nucleo scores (already 0-1, lower is better)
165    for (const fileSource of fileSources) {
166      scoredResults.push({
167        source: fileSource,
168        score: fileSource.score ?? 0.5, // Default to middle score if missing
169      })
170    }
171  
172    // Score non-file sources with Fuse.js and add them
173    if (nonFileSources.length > 0) {
174      const fuse = new Fuse(nonFileSources, {
175        includeScore: true,
176        threshold: 0.6, // Allow more matches through, we'll sort by score
177        keys: [
178          { name: 'displayText', weight: 2 },
179          { name: 'name', weight: 3 },
180          { name: 'server', weight: 1 },
181          { name: 'description', weight: 1 },
182          { name: 'agentType', weight: 3 },
183        ],
184      })
185  
186      const fuseResults = fuse.search(query, { limit: MAX_UNIFIED_SUGGESTIONS })
187      for (const result of fuseResults) {
188        scoredResults.push({
189          source: result.item,
190          score: result.score ?? 0.5,
191        })
192      }
193    }
194  
195    // Sort all results by score (lower is better) and return top results
196    scoredResults.sort((a, b) => a.score - b.score)
197  
198    return scoredResults
199      .slice(0, MAX_UNIFIED_SUGGESTIONS)
200      .map(r => r.source)
201      .map(createSuggestionFromSource)
202  }