1  import type {
  2    ContentBlockParam,
  3    ImageBlockParam,
  4    TextBlockParam,
  5  } from '@anthropic-ai/sdk/resources/index.mjs'
  6  import { getFeatureValue_CACHED_MAY_BE_STALE } from '../services/analytics/growthbook.js'
  7  import {
  8    countMessagesTokensWithAPI,
  9    roughTokenCountEstimation,
 10  } from '../services/tokenEstimation.js'
 11  import { compressImageBlock } from './imageResizer.js'
 12  import { logError } from './log.js'
 13  
 14  export const MCP_TOKEN_COUNT_THRESHOLD_FACTOR = 0.5
 15  export const IMAGE_TOKEN_ESTIMATE = 1600
 16  const DEFAULT_MAX_MCP_OUTPUT_TOKENS = 25000
 17  
 18  /**
 19   * Resolve the MCP output token cap. Precedence:
 20   *   1. MAX_MCP_OUTPUT_TOKENS env var (explicit user override)
 21   *   2. tengu_satin_quoll GrowthBook flag's `mcp_tool` key (tokens, not chars —
 22   *      unlike the other keys in that map which getPersistenceThreshold reads
 23   *      as chars; MCP has its own truncation layer upstream of that)
 24   *   3. Hardcoded default
 25   */
 26  export function getMaxMcpOutputTokens(): number {
 27    const envValue = process.env.MAX_MCP_OUTPUT_TOKENS
 28    if (envValue) {
 29      const parsed = parseInt(envValue, 10)
 30      if (Number.isFinite(parsed) && parsed > 0) {
 31        return parsed
 32      }
 33    }
 34    const overrides = getFeatureValue_CACHED_MAY_BE_STALE<Record<
 35      string,
 36      number
 37    > | null>('tengu_satin_quoll', {})
 38    const override = overrides?.['mcp_tool']
 39    if (
 40      typeof override === 'number' &&
 41      Number.isFinite(override) &&
 42      override > 0
 43    ) {
 44      return override
 45    }
 46    return DEFAULT_MAX_MCP_OUTPUT_TOKENS
 47  }
 48  
 49  export type MCPToolResult = string | ContentBlockParam[] | undefined
 50  
 51  function isTextBlock(block: ContentBlockParam): block is TextBlockParam {
 52    return block.type === 'text'
 53  }
 54  
 55  function isImageBlock(block: ContentBlockParam): block is ImageBlockParam {
 56    return block.type === 'image'
 57  }
 58  
 59  export function getContentSizeEstimate(content: MCPToolResult): number {
 60    if (!content) return 0
 61  
 62    if (typeof content === 'string') {
 63      return roughTokenCountEstimation(content)
 64    }
 65  
 66    return content.reduce((total, block) => {
 67      if (isTextBlock(block)) {
 68        return total + roughTokenCountEstimation(block.text)
 69      } else if (isImageBlock(block)) {
 70        // Estimate for image tokens
 71        return total + IMAGE_TOKEN_ESTIMATE
 72      }
 73      return total
 74    }, 0)
 75  }
 76  
 77  function getMaxMcpOutputChars(): number {
 78    return getMaxMcpOutputTokens() * 4
 79  }
 80  
 81  function getTruncationMessage(): string {
 82    return `\n\n[OUTPUT TRUNCATED - exceeded ${getMaxMcpOutputTokens()} token limit]
 83  
 84  The tool output was truncated. If this MCP server provides pagination or filtering tools, use them to retrieve specific portions of the data. If pagination is not available, inform the user that you are working with truncated output and results may be incomplete.`
 85  }
 86  
 87  function truncateString(content: string, maxChars: number): string {
 88    if (content.length <= maxChars) {
 89      return content
 90    }
 91    return content.slice(0, maxChars)
 92  }
 93  
 94  async function truncateContentBlocks(
 95    blocks: ContentBlockParam[],
 96    maxChars: number,
 97  ): Promise<ContentBlockParam[]> {
 98    const result: ContentBlockParam[] = []
 99    let currentChars = 0
100  
101    for (const block of blocks) {
102      if (isTextBlock(block)) {
103        const remainingChars = maxChars - currentChars
104        if (remainingChars <= 0) break
105  
106        if (block.text.length <= remainingChars) {
107          result.push(block)
108          currentChars += block.text.length
109        } else {
110          result.push({ type: 'text', text: block.text.slice(0, remainingChars) })
111          break
112        }
113      } else if (isImageBlock(block)) {
114        // Include images but count their estimated size
115        const imageChars = IMAGE_TOKEN_ESTIMATE * 4
116        if (currentChars + imageChars <= maxChars) {
117          result.push(block)
118          currentChars += imageChars
119        } else {
120          // Image exceeds budget - try to compress it to fit remaining space
121          const remainingChars = maxChars - currentChars
122          if (remainingChars > 0) {
123            // Convert remaining chars to bytes for compression
124            // base64 uses ~4/3 the original size, so we calculate max bytes
125            const remainingBytes = Math.floor(remainingChars * 0.75)
126            try {
127              const compressedBlock = await compressImageBlock(
128                block,
129                remainingBytes,
130              )
131              result.push(compressedBlock)
132              // Update currentChars based on compressed image size
133              if (compressedBlock.source.type === 'base64') {
134                currentChars += compressedBlock.source.data.length
135              } else {
136                currentChars += imageChars
137              }
138            } catch {
139              // If compression fails, skip the image
140            }
141          }
142        }
143      } else {
144        result.push(block)
145      }
146    }
147  
148    return result
149  }
150  
151  export async function mcpContentNeedsTruncation(
152    content: MCPToolResult,
153  ): Promise<boolean> {
154    if (!content) return false
155  
156    // Use size check as a heuristic to avoid unnecessary token counting API calls
157    const contentSizeEstimate = getContentSizeEstimate(content)
158    if (
159      contentSizeEstimate <=
160      getMaxMcpOutputTokens() * MCP_TOKEN_COUNT_THRESHOLD_FACTOR
161    ) {
162      return false
163    }
164  
165    try {
166      const messages =
167        typeof content === 'string'
168          ? [{ role: 'user' as const, content }]
169          : [{ role: 'user' as const, content }]
170  
171      const tokenCount = await countMessagesTokensWithAPI(messages, [])
172      return !!(tokenCount && tokenCount > getMaxMcpOutputTokens())
173    } catch (error) {
174      logError(error)
175      // Assume no truncation needed on error
176      return false
177    }
178  }
179  
180  export async function truncateMcpContent(
181    content: MCPToolResult,
182  ): Promise<MCPToolResult> {
183    if (!content) return content
184  
185    const maxChars = getMaxMcpOutputChars()
186    const truncationMsg = getTruncationMessage()
187  
188    if (typeof content === 'string') {
189      return truncateString(content, maxChars) + truncationMsg
190    } else {
191      const truncatedBlocks = await truncateContentBlocks(
192        content as ContentBlockParam[],
193        maxChars,
194      )
195      truncatedBlocks.push({ type: 'text', text: truncationMsg })
196      return truncatedBlocks
197    }
198  }
199  
200  export async function truncateMcpContentIfNeeded(
201    content: MCPToolResult,
202  ): Promise<MCPToolResult> {
203    if (!(await mcpContentNeedsTruncation(content))) {
204      return content
205    }
206  
207    return await truncateMcpContent(content)
208  }