1  import { roughTokenCountEstimation } from '../services/tokenEstimation.js'
  2  import type { Tool, ToolPermissionContext } from '../Tool.js'
  3  import type { AgentDefinitionsResult } from '../tools/AgentTool/loadAgentsDir.js'
  4  import { countMcpToolTokens } from './analyzeContext.js'
  5  import {
  6    getLargeMemoryFiles,
  7    getMemoryFiles,
  8    MAX_MEMORY_CHARACTER_COUNT,
  9  } from './claudemd.js'
 10  import { getMainLoopModel } from './model/model.js'
 11  import { permissionRuleValueToString } from './permissions/permissionRuleParser.js'
 12  import { detectUnreachableRules } from './permissions/shadowedRuleDetection.js'
 13  import { SandboxManager } from './sandbox/sandbox-adapter.js'
 14  import {
 15    AGENT_DESCRIPTIONS_THRESHOLD,
 16    getAgentDescriptionsTotalTokens,
 17  } from './statusNoticeHelpers.js'
 18  import { plural } from './stringUtils.js'
 19  
 20  // Thresholds (matching status notices and existing patterns)
 21  const MCP_TOOLS_THRESHOLD = 25_000 // 15k tokens
 22  
 23  export type ContextWarning = {
 24    type:
 25      | 'claudemd_files'
 26      | 'agent_descriptions'
 27      | 'mcp_tools'
 28      | 'unreachable_rules'
 29    severity: 'warning' | 'error'
 30    message: string
 31    details: string[]
 32    currentValue: number
 33    threshold: number
 34  }
 35  
 36  export type ContextWarnings = {
 37    claudeMdWarning: ContextWarning | null
 38    agentWarning: ContextWarning | null
 39    mcpWarning: ContextWarning | null
 40    unreachableRulesWarning: ContextWarning | null
 41  }
 42  
 43  async function checkClaudeMdFiles(): Promise<ContextWarning | null> {
 44    const largeFiles = getLargeMemoryFiles(await getMemoryFiles())
 45  
 46    // This already filters for files > 40k chars each
 47    if (largeFiles.length === 0) {
 48      return null
 49    }
 50  
 51    const details = largeFiles
 52      .sort((a, b) => b.content.length - a.content.length)
 53      .map(file => `${file.path}: ${file.content.length.toLocaleString()} chars`)
 54  
 55    const message =
 56      largeFiles.length === 1
 57        ? `Large CLAUDE.md file detected (${largeFiles[0]!.content.length.toLocaleString()} chars > ${MAX_MEMORY_CHARACTER_COUNT.toLocaleString()})`
 58        : `${largeFiles.length} large CLAUDE.md files detected (each > ${MAX_MEMORY_CHARACTER_COUNT.toLocaleString()} chars)`
 59  
 60    return {
 61      type: 'claudemd_files',
 62      severity: 'warning',
 63      message,
 64      details,
 65      currentValue: largeFiles.length, // Number of files exceeding threshold
 66      threshold: MAX_MEMORY_CHARACTER_COUNT,
 67    }
 68  }
 69  
 70  /**
 71   * Check agent descriptions token count
 72   */
 73  async function checkAgentDescriptions(
 74    agentInfo: AgentDefinitionsResult | null,
 75  ): Promise<ContextWarning | null> {
 76    if (!agentInfo) {
 77      return null
 78    }
 79  
 80    const totalTokens = getAgentDescriptionsTotalTokens(agentInfo)
 81  
 82    if (totalTokens <= AGENT_DESCRIPTIONS_THRESHOLD) {
 83      return null
 84    }
 85  
 86    // Calculate tokens for each agent
 87    const agentTokens = agentInfo.activeAgents
 88      .filter(a => a.source !== 'built-in')
 89      .map(agent => {
 90        const description = `${agent.agentType}: ${agent.whenToUse}`
 91        return {
 92          name: agent.agentType,
 93          tokens: roughTokenCountEstimation(description),
 94        }
 95      })
 96      .sort((a, b) => b.tokens - a.tokens)
 97  
 98    const details = agentTokens
 99      .slice(0, 5)
100      .map(agent => `${agent.name}: ~${agent.tokens.toLocaleString()} tokens`)
101  
102    if (agentTokens.length > 5) {
103      details.push(`(${agentTokens.length - 5} more custom agents)`)
104    }
105  
106    return {
107      type: 'agent_descriptions',
108      severity: 'warning',
109      message: `Large agent descriptions (~${totalTokens.toLocaleString()} tokens > ${AGENT_DESCRIPTIONS_THRESHOLD.toLocaleString()})`,
110      details,
111      currentValue: totalTokens,
112      threshold: AGENT_DESCRIPTIONS_THRESHOLD,
113    }
114  }
115  
116  /**
117   * Check MCP tools token count
118   */
119  async function checkMcpTools(
120    tools: Tool[],
121    getToolPermissionContext: () => Promise<ToolPermissionContext>,
122    agentInfo: AgentDefinitionsResult | null,
123  ): Promise<ContextWarning | null> {
124    const mcpTools = tools.filter(tool => tool.isMcp)
125  
126    // Note: MCP tools are loaded asynchronously and may not be available
127    // when doctor command runs, as it executes before MCP connections are established
128    if (mcpTools.length === 0) {
129      return null
130    }
131  
132    try {
133      // Use the existing countMcpToolTokens function from analyzeContext
134      const model = getMainLoopModel()
135      const { mcpToolTokens, mcpToolDetails } = await countMcpToolTokens(
136        tools,
137        getToolPermissionContext,
138        agentInfo,
139        model,
140      )
141  
142      if (mcpToolTokens <= MCP_TOOLS_THRESHOLD) {
143        return null
144      }
145  
146      // Group tools by server
147      const toolsByServer = new Map<string, { count: number; tokens: number }>()
148  
149      for (const tool of mcpToolDetails) {
150        // Extract server name from tool name (format: mcp__servername__toolname)
151        const parts = tool.name.split('__')
152        const serverName = parts[1] || 'unknown'
153  
154        const current = toolsByServer.get(serverName) || { count: 0, tokens: 0 }
155        toolsByServer.set(serverName, {
156          count: current.count + 1,
157          tokens: current.tokens + tool.tokens,
158        })
159      }
160  
161      // Sort servers by token count
162      const sortedServers = Array.from(toolsByServer.entries()).sort(
163        (a, b) => b[1].tokens - a[1].tokens,
164      )
165  
166      const details = sortedServers
167        .slice(0, 5)
168        .map(
169          ([name, info]) =>
170            `${name}: ${info.count} tools (~${info.tokens.toLocaleString()} tokens)`,
171        )
172  
173      if (sortedServers.length > 5) {
174        details.push(`(${sortedServers.length - 5} more servers)`)
175      }
176  
177      return {
178        type: 'mcp_tools',
179        severity: 'warning',
180        message: `Large MCP tools context (~${mcpToolTokens.toLocaleString()} tokens > ${MCP_TOOLS_THRESHOLD.toLocaleString()})`,
181        details,
182        currentValue: mcpToolTokens,
183        threshold: MCP_TOOLS_THRESHOLD,
184      }
185    } catch (_error) {
186      // If token counting fails, fall back to character-based estimation
187      const estimatedTokens = mcpTools.reduce((total, tool) => {
188        const chars = (tool.name?.length || 0) + tool.description.length
189        return total + roughTokenCountEstimation(chars.toString())
190      }, 0)
191  
192      if (estimatedTokens <= MCP_TOOLS_THRESHOLD) {
193        return null
194      }
195  
196      return {
197        type: 'mcp_tools',
198        severity: 'warning',
199        message: `Large MCP tools context (~${estimatedTokens.toLocaleString()} tokens estimated > ${MCP_TOOLS_THRESHOLD.toLocaleString()})`,
200        details: [
201          `${mcpTools.length} MCP tools detected (token count estimated)`,
202        ],
203        currentValue: estimatedTokens,
204        threshold: MCP_TOOLS_THRESHOLD,
205      }
206    }
207  }
208  
209  /**
210   * Check for unreachable permission rules (e.g., specific allow rules shadowed by tool-wide ask rules)
211   */
212  async function checkUnreachableRules(
213    getToolPermissionContext: () => Promise<ToolPermissionContext>,
214  ): Promise<ContextWarning | null> {
215    const context = await getToolPermissionContext()
216    const sandboxAutoAllowEnabled =
217      SandboxManager.isSandboxingEnabled() &&
218      SandboxManager.isAutoAllowBashIfSandboxedEnabled()
219  
220    const unreachable = detectUnreachableRules(context, {
221      sandboxAutoAllowEnabled,
222    })
223  
224    if (unreachable.length === 0) {
225      return null
226    }
227  
228    const details = unreachable.flatMap(r => [
229      `${permissionRuleValueToString(r.rule.ruleValue)}: ${r.reason}`,
230      `  Fix: ${r.fix}`,
231    ])
232  
233    return {
234      type: 'unreachable_rules',
235      severity: 'warning',
236      message: `${unreachable.length} ${plural(unreachable.length, 'unreachable permission rule')} detected`,
237      details,
238      currentValue: unreachable.length,
239      threshold: 0,
240    }
241  }
242  
243  /**
244   * Check all context warnings for the doctor command
245   */
246  export async function checkContextWarnings(
247    tools: Tool[],
248    agentInfo: AgentDefinitionsResult | null,
249    getToolPermissionContext: () => Promise<ToolPermissionContext>,
250  ): Promise<ContextWarnings> {
251    const [claudeMdWarning, agentWarning, mcpWarning, unreachableRulesWarning] =
252      await Promise.all([
253        checkClaudeMdFiles(),
254        checkAgentDescriptions(agentInfo),
255        checkMcpTools(tools, getToolPermissionContext, agentInfo),
256        checkUnreachableRules(getToolPermissionContext),
257      ])
258  
259    return {
260      claudeMdWarning,
261      agentWarning,
262      mcpWarning,
263      unreachableRulesWarning,
264    }
265  }