1  import { FILE_EDIT_TOOL_NAME } from 'src/tools/FileEditTool/constants.js'
  2  import { FILE_READ_TOOL_NAME } from 'src/tools/FileReadTool/prompt.js'
  3  import { FILE_WRITE_TOOL_NAME } from 'src/tools/FileWriteTool/prompt.js'
  4  import { GLOB_TOOL_NAME } from 'src/tools/GlobTool/prompt.js'
  5  import { GREP_TOOL_NAME } from 'src/tools/GrepTool/prompt.js'
  6  import { NOTEBOOK_EDIT_TOOL_NAME } from 'src/tools/NotebookEditTool/constants.js'
  7  import { WEB_FETCH_TOOL_NAME } from 'src/tools/WebFetchTool/prompt.js'
  8  import { WEB_SEARCH_TOOL_NAME } from 'src/tools/WebSearchTool/prompt.js'
  9  import { SHELL_TOOL_NAMES } from 'src/utils/shell/shellToolUtils.js'
 10  import { isEnvTruthy } from '../../utils/envUtils.js'
 11  
 12  // docs: https://docs.google.com/document/d/1oCT4evvWTh3P6z-kcfNQwWTCxAhkoFndSaNS9Gm40uw/edit?tab=t.0
 13  
 14  // Default values for context management strategies
 15  // Match client-side microcompact token values
 16  const DEFAULT_MAX_INPUT_TOKENS = 180_000 // Typical warning threshold
 17  const DEFAULT_TARGET_INPUT_TOKENS = 40_000 // Keep last 40k tokens like client-side
 18  
 19  const TOOLS_CLEARABLE_RESULTS = [
 20    ...SHELL_TOOL_NAMES,
 21    GLOB_TOOL_NAME,
 22    GREP_TOOL_NAME,
 23    FILE_READ_TOOL_NAME,
 24    WEB_FETCH_TOOL_NAME,
 25    WEB_SEARCH_TOOL_NAME,
 26  ]
 27  
 28  const TOOLS_CLEARABLE_USES = [
 29    FILE_EDIT_TOOL_NAME,
 30    FILE_WRITE_TOOL_NAME,
 31    NOTEBOOK_EDIT_TOOL_NAME,
 32  ]
 33  
 34  // Context management strategy types matching API documentation
 35  export type ContextEditStrategy =
 36    | {
 37        type: 'clear_tool_uses_20250919'
 38        trigger?: {
 39          type: 'input_tokens'
 40          value: number
 41        }
 42        keep?: {
 43          type: 'tool_uses'
 44          value: number
 45        }
 46        clear_tool_inputs?: boolean | string[]
 47        exclude_tools?: string[]
 48        clear_at_least?: {
 49          type: 'input_tokens'
 50          value: number
 51        }
 52      }
 53    | {
 54        type: 'clear_thinking_20251015'
 55        keep: { type: 'thinking_turns'; value: number } | 'all'
 56      }
 57  
 58  // Context management configuration wrapper
 59  export type ContextManagementConfig = {
 60    edits: ContextEditStrategy[]
 61  }
 62  
 63  // API-based microcompact implementation that uses native context management
 64  export function getAPIContextManagement(options?: {
 65    hasThinking?: boolean
 66    isRedactThinkingActive?: boolean
 67    clearAllThinking?: boolean
 68  }): ContextManagementConfig | undefined {
 69    const {
 70      hasThinking = false,
 71      isRedactThinkingActive = false,
 72      clearAllThinking = false,
 73    } = options ?? {}
 74  
 75    const strategies: ContextEditStrategy[] = []
 76  
 77    // Preserve thinking blocks in previous assistant turns. Skip when
 78    // redact-thinking is active — redacted blocks have no model-visible content.
 79    // When clearAllThinking is set (>1h idle = cache miss), keep only the last
 80    // thinking turn — the API schema requires value >= 1, and omitting the edit
 81    // falls back to the model-policy default (often "all"), which wouldn't clear.
 82    if (hasThinking && !isRedactThinkingActive) {
 83      strategies.push({
 84        type: 'clear_thinking_20251015',
 85        keep: clearAllThinking ? { type: 'thinking_turns', value: 1 } : 'all',
 86      })
 87    }
 88  
 89    // Tool clearing strategies are ant-only
 90    if (process.env.USER_TYPE !== 'ant') {
 91      return strategies.length > 0 ? { edits: strategies } : undefined
 92    }
 93  
 94    const useClearToolResults = isEnvTruthy(
 95      process.env.USE_API_CLEAR_TOOL_RESULTS,
 96    )
 97    const useClearToolUses = isEnvTruthy(process.env.USE_API_CLEAR_TOOL_USES)
 98  
 99    // If no tool clearing strategy is enabled, return early
100    if (!useClearToolResults && !useClearToolUses) {
101      return strategies.length > 0 ? { edits: strategies } : undefined
102    }
103  
104    if (useClearToolResults) {
105      const triggerThreshold = process.env.API_MAX_INPUT_TOKENS
106        ? parseInt(process.env.API_MAX_INPUT_TOKENS)
107        : DEFAULT_MAX_INPUT_TOKENS
108      const keepTarget = process.env.API_TARGET_INPUT_TOKENS
109        ? parseInt(process.env.API_TARGET_INPUT_TOKENS)
110        : DEFAULT_TARGET_INPUT_TOKENS
111  
112      const strategy: ContextEditStrategy = {
113        type: 'clear_tool_uses_20250919',
114        trigger: {
115          type: 'input_tokens',
116          value: triggerThreshold,
117        },
118        clear_at_least: {
119          type: 'input_tokens',
120          value: triggerThreshold - keepTarget,
121        },
122        clear_tool_inputs: TOOLS_CLEARABLE_RESULTS,
123      }
124  
125      strategies.push(strategy)
126    }
127  
128    if (useClearToolUses) {
129      const triggerThreshold = process.env.API_MAX_INPUT_TOKENS
130        ? parseInt(process.env.API_MAX_INPUT_TOKENS)
131        : DEFAULT_MAX_INPUT_TOKENS
132      const keepTarget = process.env.API_TARGET_INPUT_TOKENS
133        ? parseInt(process.env.API_TARGET_INPUT_TOKENS)
134        : DEFAULT_TARGET_INPUT_TOKENS
135  
136      const strategy: ContextEditStrategy = {
137        type: 'clear_tool_uses_20250919',
138        trigger: {
139          type: 'input_tokens',
140          value: triggerThreshold,
141        },
142        clear_at_least: {
143          type: 'input_tokens',
144          value: triggerThreshold - keepTarget,
145        },
146        exclude_tools: TOOLS_CLEARABLE_USES,
147      }
148  
149      strategies.push(strategy)
150    }
151  
152    return strategies.length > 0 ? { edits: strategies } : undefined
153  }