Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 104 additions & 14 deletions src/services/mcp/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,10 @@ import {
} from '../../utils/errors.js'
import { getMCPUserAgent } from '../../utils/http.js'
import { maybeNotifyIDEConnected } from '../../utils/ide.js'
import { maybeResizeAndDownsampleImageBuffer } from '../../utils/imageResizer.js'
import {
type ImageLimits,
maybeResizeAndDownsampleImageBuffer,
} from '../../utils/imageResizer.js'
import { logMCPDebug, logMCPError } from '../../utils/log.js'
import {
getBinaryBlobSavedMessage,
Expand Down Expand Up @@ -102,6 +105,7 @@ import {
isPersistError,
persistToolResult,
} from '../../utils/toolResultStorage.js'
import { checkStatsigFeatureGate_CACHED_MAY_BE_STALE } from '../analytics/growthbook.js'
import {
type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
logEvent,
Expand Down Expand Up @@ -2486,15 +2490,23 @@ export function prefetchAllMcpResources(
export async function transformResultContent(
resultContent: PromptMessage['content'],
serverName: string,
limits?: ImageLimits,
includeMeta = false,
): Promise<Array<ContentBlockParam>> {
switch (resultContent.type) {
case 'text':
return [
{
type: 'text',
text: resultContent.text,
},
]
case 'text': {
const block: ContentBlockParam = {
type: 'text',
text: resultContent.text,
}
if (includeMeta) {
const meta = resultContent._meta
if (meta) {
;(block as { _meta?: unknown })._meta = meta
}
}
return [block]
Comment on lines +2497 to +2508
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Preserve annotations before the single-text unwrap path runs.

unwrapSingleTextBlock() explicitly refuses annotated text blocks, but this transformer only carries _meta. A single annotated MCP text result is therefore converted into a bare text block here and later qualifies for plain-text unwrapping, which drops the annotations on the persisted large-output path.

Suggested fix
     case 'text': {
       const block: ContentBlockParam = {
         type: 'text',
         text: resultContent.text,
       }
+      if ('annotations' in resultContent && resultContent.annotations) {
+        ;(block as { annotations?: unknown }).annotations =
+          resultContent.annotations
+      }
       if (includeMeta) {
         const meta = resultContent._meta
         if (meta) {
           ;(block as { _meta?: unknown })._meta = meta
         }
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/services/mcp/client.ts` around lines 2497 - 2508, The text-case
transformer currently only copies `_meta` and thus strips `annotations` before
`unwrapSingleTextBlock()` runs; change the `case 'text'` path (where `block:
ContentBlockParam` is constructed from `resultContent.text`) to also preserve
any `resultContent.annotations` by assigning them onto the block (e.g. if
(resultContent.annotations) (block as { annotations?: unknown }).annotations =
resultContent.annotations), so annotated text blocks are not converted into bare
text and lost by the unwrapping logic.

}
case 'audio': {
const audioData = resultContent as {
type: 'audio'
Expand All @@ -2516,6 +2528,7 @@ export async function transformResultContent(
imageBuffer,
imageBuffer.length,
ext,
limits,
)
return [
{
Expand Down Expand Up @@ -2551,6 +2564,7 @@ export async function transformResultContent(
imageBuffer,
imageBuffer.length,
ext,
limits,
)
const content: MessageParam['content'] = []
if (prefix) {
Expand Down Expand Up @@ -2671,6 +2685,7 @@ export async function transformMCPResult(
result: unknown,
tool: string, // Tool name for validation (e.g., "search")
name: string, // Server name for transformation (e.g., "slack")
limits?: ImageLimits, // Image processing limits, plumbed to transformResultContent
): Promise<TransformedMCPResult> {
if (result && typeof result === 'object') {
if ('toolResult' in result) {
Expand All @@ -2694,7 +2709,9 @@ export async function transformMCPResult(
if ('content' in result && Array.isArray(result.content)) {
const transformedContent = (
await Promise.all(
result.content.map(item => transformResultContent(item, name)),
result.content.map(item =>
transformResultContent(item, name, limits, true),
),
)
).flat()
return {
Expand Down Expand Up @@ -2729,15 +2746,28 @@ export async function processMCPResult(
result: unknown,
tool: string, // Tool name for validation (e.g., "search")
name: string, // Server name for IDE check and transformation (e.g., "slack")
limits?: ImageLimits, // Image processing limits, plumbed to transformMCPResult
skipLargeOutput = false, // If true, skip large-output handling for non-image content
): Promise<MCPToolResult> {
const { content, type, schema } = await transformMCPResult(result, tool, name)
const { content, type, schema } = await transformMCPResult(
result,
tool,
name,
limits,
)

// IDE tools are not going to the model directly, so we don't need to
// handle large output.
if (name === 'ide') {
return content
}

// Caller opted out of large-output handling (e.g., result already truncated
// upstream); only continue if the content has images that may need handling.
if (skipLargeOutput && !contentContainsImages(content)) {
return content
}

// Check if content needs truncation (i.e., is too large)
if (!(await mcpContentNeedsTruncation(content))) {
return content
Expand Down Expand Up @@ -2775,9 +2805,15 @@ export async function processMCPResult(
// Generate a unique ID for the persisted file (server__tool-timestamp)
const timestamp = Date.now()
const persistId = `mcp-${normalizeNameForMCP(name)}-${normalizeNameForMCP(tool)}-${timestamp}`
// Convert to string for persistence (persistToolResult expects string or specific block types)
// When the large-string format gate is on, unwrap a single bare text block
// (no annotations, no _meta) into raw text so the model gets plain text in
// the persisted file instead of a JSON-wrapped block. The `_meta` check is
// why transformResultContent preserves _meta on text blocks.
const unwrappedText = unwrapSingleTextBlock(content)
const contentStr =
typeof content === 'string' ? content : jsonStringify(content, null, 2)
typeof content === 'string'
? content
: (unwrappedText ?? jsonStringify(content, null, 2))
const persistResult = await persistToolResult(contentStr, persistId)

if (isPersistError(persistResult)) {
Expand All @@ -2798,14 +2834,50 @@ export async function processMCPResult(
persistedSizeChars: persistResult.originalSize,
} as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS)

const formatDescription = getFormatDescription(type, schema)
const formatDescription =
unwrappedText !== undefined
? getFormatDescription('toolResult')
: getFormatDescription(type, schema)
return getLargeOutputInstructions(
persistResult.filepath,
persistResult.originalSize,
formatDescription,
)
}

/**
* Returns true when the large-string output format is enabled (matches
* binary's `sf8()`). When enabled, processMCPResult unwraps a single bare
* text block to raw text for persistence instead of JSON-wrapping it.
*
* Gating sources, in order:
* 1. MCP_TRUNCATION_PROMPT_OVERRIDE env var (anything except "legacy" enables)
* 2. Statsig gate `tengu_mcp_subagent_prompt`
*/
function isLargeStringFormatEnabled(): boolean {
const override = process.env.MCP_TRUNCATION_PROMPT_OVERRIDE
if (override) return override !== 'legacy'
return checkStatsigFeatureGate_CACHED_MAY_BE_STALE(
'tengu_mcp_subagent_prompt',
)
}

/**
* Unwraps a single bare text content block to its raw text when the
* large-string format gate is on. Returns undefined when the gate is off,
* the content is not an array, the array doesn't contain exactly one text
* block, or the block carries annotations or _meta. Matches binary mg5's
* `M=...` computation.
*/
function unwrapSingleTextBlock(content: MCPToolResult): string | undefined {
if (!isLargeStringFormatEnabled()) return undefined
if (!Array.isArray(content) || content.length !== 1) return undefined
const block = content[0]
if (!block || block.type !== 'text') return undefined
if ('annotations' in block || '_meta' in block) return undefined
return block.text
}

/**
* Call an MCP tool, handling UrlElicitationRequiredError (-32042) by
* displaying the URL elicitation to the user, waiting for the completion
Expand All @@ -2827,6 +2899,8 @@ export async function callMCPToolWithUrlElicitationRetry({
signal,
setAppState,
onProgress,
imageLimits,
hasResultSizeAnnotation = false,
callToolFn = callMCPTool,
handleElicitation,
}: {
Expand All @@ -2838,6 +2912,8 @@ export async function callMCPToolWithUrlElicitationRetry({
signal: AbortSignal
setAppState: (f: (prev: AppState) => AppState) => void
onProgress?: (data: MCPProgress) => void
imageLimits?: ImageLimits
hasResultSizeAnnotation?: boolean
/** Injectable for testing. Defaults to callMCPTool. */
callToolFn?: (opts: {
client: ConnectedMCPServer
Expand All @@ -2846,6 +2922,8 @@ export async function callMCPToolWithUrlElicitationRetry({
meta?: Record<string, unknown>
signal: AbortSignal
onProgress?: (data: MCPProgress) => void
imageLimits?: ImageLimits
hasResultSizeAnnotation?: boolean
}) => Promise<MCPToolCallResult>
/** Handler for URL elicitations when no hook handles them.
* In print/SDK mode, delegates to structuredIO. In REPL, falls back to queue. */
Expand All @@ -2865,6 +2943,8 @@ export async function callMCPToolWithUrlElicitationRetry({
meta,
signal,
onProgress,
imageLimits,
hasResultSizeAnnotation,
})
} catch (error) {
// The MCP SDK's Protocol creates plain McpError (not UrlElicitationRequiredError)
Expand Down Expand Up @@ -3041,13 +3121,17 @@ async function callMCPTool({
meta,
signal,
onProgress,
imageLimits,
hasResultSizeAnnotation = false,
}: {
client: ConnectedMCPServer
tool: string
args: Record<string, unknown>
meta?: Record<string, unknown>
signal: AbortSignal
onProgress?: (data: MCPProgress) => void
imageLimits?: ImageLimits
hasResultSizeAnnotation?: boolean
}): Promise<{
content: MCPToolResult
_meta?: Record<string, unknown>
Expand Down Expand Up @@ -3176,7 +3260,13 @@ async function callMCPTool({
})
}

const content = await processMCPResult(result, tool, name)
const content = await processMCPResult(
result,
tool,
name,
imageLimits,
hasResultSizeAnnotation,
)
return {
content,
_meta: result._meta as Record<string, unknown> | undefined,
Expand Down
Loading