/ utils / queryHelpers.ts
queryHelpers.ts
  1  import type { ToolUseBlock } from '@anthropic-ai/sdk/resources/index.mjs'
  2  import last from 'lodash-es/last.js'
  3  import {
  4    getSessionId,
  5    isSessionPersistenceDisabled,
  6  } from 'src/bootstrap/state.js'
  7  import type { SDKMessage } from 'src/entrypoints/agentSdkTypes.js'
  8  import type { CanUseToolFn } from '../hooks/useCanUseTool.js'
  9  import { runTools } from '../services/tools/toolOrchestration.js'
 10  import { findToolByName, type Tool, type Tools } from '../Tool.js'
 11  import { BASH_TOOL_NAME } from '../tools/BashTool/toolName.js'
 12  import { FILE_EDIT_TOOL_NAME } from '../tools/FileEditTool/constants.js'
 13  import type { Input as FileReadInput } from '../tools/FileReadTool/FileReadTool.js'
 14  import {
 15    FILE_READ_TOOL_NAME,
 16    FILE_UNCHANGED_STUB,
 17  } from '../tools/FileReadTool/prompt.js'
 18  import { FILE_WRITE_TOOL_NAME } from '../tools/FileWriteTool/prompt.js'
 19  import type { Message } from '../types/message.js'
 20  import type { OrphanedPermission } from '../types/textInputTypes.js'
 21  import { logForDebugging } from './debug.js'
 22  import { isEnvTruthy } from './envUtils.js'
 23  import { isFsInaccessible } from './errors.js'
 24  import { getFileModificationTime, stripLineNumberPrefix } from './file.js'
 25  import { readFileSyncWithMetadata } from './fileRead.js'
 26  import {
 27    createFileStateCacheWithSizeLimit,
 28    type FileStateCache,
 29  } from './fileStateCache.js'
 30  import { isNotEmptyMessage, normalizeMessages } from './messages.js'
 31  import { expandPath } from './path.js'
 32  import type {
 33    inputSchema as permissionToolInputSchema,
 34    outputSchema as permissionToolOutputSchema,
 35  } from './permissions/PermissionPromptToolResultSchema.js'
 36  import type { ProcessUserInputContext } from './processUserInput/processUserInput.js'
 37  import { recordTranscript } from './sessionStorage.js'
 38  
 39  export type PermissionPromptTool = Tool<
 40    ReturnType<typeof permissionToolInputSchema>,
 41    ReturnType<typeof permissionToolOutputSchema>
 42  >
 43  
 44  // Small cache size for ask operations which typically access few files
 45  // during permission prompts or limited tool operations
 46  const ASK_READ_FILE_STATE_CACHE_SIZE = 10
 47  
 48  /**
 49   * Checks if the result should be considered successful based on the last message.
 50   * Returns true if:
 51   * - Last message is assistant with text/thinking content
 52   * - Last message is user with only tool_result blocks
 53   * - Last message is the user prompt but the API completed with end_turn
 54   *   (model chose to emit no content blocks)
 55   */
 56  export function isResultSuccessful(
 57    message: Message | undefined,
 58    stopReason: string | null = null,
 59  ): message is Message {
 60    if (!message) return false
 61  
 62    if (message.type === 'assistant') {
 63      const lastContent = last(message.message.content)
 64      return (
 65        lastContent?.type === 'text' ||
 66        lastContent?.type === 'thinking' ||
 67        lastContent?.type === 'redacted_thinking'
 68      )
 69    }
 70  
 71    if (message.type === 'user') {
 72      // Check if all content blocks are tool_result type
 73      const content = message.message.content
 74      if (
 75        Array.isArray(content) &&
 76        content.length > 0 &&
 77        content.every(block => 'type' in block && block.type === 'tool_result')
 78      ) {
 79        return true
 80      }
 81    }
 82  
 83    // Carve-out: API completed (message_delta set stop_reason) but yielded
 84    // no assistant content — last(messages) is still this turn's prompt.
 85    // claude.ts:2026 recognizes end_turn-with-zero-content-blocks as
 86    // legitimate and passes through without throwing. Observed on
 87    // task_notification drain turns: model returns stop_reason=end_turn,
 88    // outputTokens=4, textContentLength=0 — it saw the subagent result
 89    // and decided nothing needed saying. Without this, QueryEngine emits
 90    // error_during_execution with errors[] = the entire process's
 91    // accumulated logError() buffer. Covers both string-content and
 92    // text-block-content user prompts, and any other non-passing shape.
 93    return stopReason === 'end_turn'
 94  }
 95  
 96  // Track last sent time for tool progress messages per tool use ID
 97  // Keep only the last 100 entries to prevent unbounded growth
 98  const MAX_TOOL_PROGRESS_TRACKING_ENTRIES = 100
 99  const TOOL_PROGRESS_THROTTLE_MS = 30000
100  const toolProgressLastSentTime = new Map<string, number>()
101  
102  export function* normalizeMessage(message: Message): Generator<SDKMessage> {
103    switch (message.type) {
104      case 'assistant':
105        for (const _ of normalizeMessages([message])) {
106          // Skip empty messages (e.g., "(no content)") that shouldn't be output to SDK
107          if (!isNotEmptyMessage(_)) {
108            continue
109          }
110          yield {
111            type: 'assistant',
112            message: _.message,
113            parent_tool_use_id: null,
114            session_id: getSessionId(),
115            uuid: _.uuid,
116            error: _.error,
117          }
118        }
119        return
120      case 'progress':
121        if (
122          message.data.type === 'agent_progress' ||
123          message.data.type === 'skill_progress'
124        ) {
125          for (const _ of normalizeMessages([message.data.message])) {
126            switch (_.type) {
127              case 'assistant':
128                // Skip empty messages (e.g., "(no content)") that shouldn't be output to SDK
129                if (!isNotEmptyMessage(_)) {
130                  break
131                }
132                yield {
133                  type: 'assistant',
134                  message: _.message,
135                  parent_tool_use_id: message.parentToolUseID,
136                  session_id: getSessionId(),
137                  uuid: _.uuid,
138                  error: _.error,
139                }
140                break
141              case 'user':
142                yield {
143                  type: 'user',
144                  message: _.message,
145                  parent_tool_use_id: message.parentToolUseID,
146                  session_id: getSessionId(),
147                  uuid: _.uuid,
148                  timestamp: _.timestamp,
149                  isSynthetic: _.isMeta || _.isVisibleInTranscriptOnly,
150                  tool_use_result: _.mcpMeta
151                    ? { content: _.toolUseResult, ..._.mcpMeta }
152                    : _.toolUseResult,
153                }
154                break
155            }
156          }
157        } else if (
158          message.data.type === 'bash_progress' ||
159          message.data.type === 'powershell_progress'
160        ) {
161          // Filter bash progress to send only one per minute
162          // Only emit for Claude Code Remote for now
163          if (
164            !isEnvTruthy(process.env.CLAUDE_CODE_REMOTE) &&
165            !process.env.CLAUDE_CODE_CONTAINER_ID
166          ) {
167            break
168          }
169  
170          // Use parentToolUseID as the key since toolUseID changes for each progress message
171          const trackingKey = message.parentToolUseID
172          const now = Date.now()
173          const lastSent = toolProgressLastSentTime.get(trackingKey) || 0
174          const timeSinceLastSent = now - lastSent
175  
176          // Send if at least 30 seconds have passed since last update
177          if (timeSinceLastSent >= TOOL_PROGRESS_THROTTLE_MS) {
178            // Remove oldest entry if we're at capacity (LRU eviction)
179            if (
180              toolProgressLastSentTime.size >= MAX_TOOL_PROGRESS_TRACKING_ENTRIES
181            ) {
182              const firstKey = toolProgressLastSentTime.keys().next().value
183              if (firstKey !== undefined) {
184                toolProgressLastSentTime.delete(firstKey)
185              }
186            }
187  
188            toolProgressLastSentTime.set(trackingKey, now)
189            yield {
190              type: 'tool_progress',
191              tool_use_id: message.toolUseID,
192              tool_name:
193                message.data.type === 'bash_progress' ? 'Bash' : 'PowerShell',
194              parent_tool_use_id: message.parentToolUseID,
195              elapsed_time_seconds: message.data.elapsedTimeSeconds,
196              task_id: message.data.taskId,
197              session_id: getSessionId(),
198              uuid: message.uuid,
199            }
200          }
201        }
202        break
203      case 'user':
204        for (const _ of normalizeMessages([message])) {
205          yield {
206            type: 'user',
207            message: _.message,
208            parent_tool_use_id: null,
209            session_id: getSessionId(),
210            uuid: _.uuid,
211            timestamp: _.timestamp,
212            isSynthetic: _.isMeta || _.isVisibleInTranscriptOnly,
213            tool_use_result: _.mcpMeta
214              ? { content: _.toolUseResult, ..._.mcpMeta }
215              : _.toolUseResult,
216          }
217        }
218        return
219      default:
220      // yield nothing
221    }
222  }
223  
224  export async function* handleOrphanedPermission(
225    orphanedPermission: OrphanedPermission,
226    tools: Tools,
227    mutableMessages: Message[],
228    processUserInputContext: ProcessUserInputContext,
229  ): AsyncGenerator<SDKMessage, void, unknown> {
230    const persistSession = !isSessionPersistenceDisabled()
231    const { permissionResult, assistantMessage } = orphanedPermission
232    const { toolUseID } = permissionResult
233  
234    if (!toolUseID) {
235      return
236    }
237  
238    const content = assistantMessage.message.content
239    let toolUseBlock: ToolUseBlock | undefined
240    if (Array.isArray(content)) {
241      for (const block of content) {
242        if (block.type === 'tool_use' && block.id === toolUseID) {
243          toolUseBlock = block as ToolUseBlock
244          break
245        }
246      }
247    }
248  
249    if (!toolUseBlock) {
250      return
251    }
252  
253    const toolName = toolUseBlock.name
254    const toolInput = toolUseBlock.input
255  
256    const toolDefinition = findToolByName(tools, toolName)
257    if (!toolDefinition) {
258      return
259    }
260  
261    // Create ToolUseBlock with the updated input if permission was allowed
262    let finalInput = toolInput
263    if (permissionResult.behavior === 'allow') {
264      if (permissionResult.updatedInput !== undefined) {
265        finalInput = permissionResult.updatedInput
266      } else {
267        logForDebugging(
268          `Orphaned permission for ${toolName}: updatedInput is undefined, falling back to original tool input`,
269          { level: 'warn' },
270        )
271      }
272    }
273    const finalToolUseBlock: ToolUseBlock = {
274      ...toolUseBlock,
275      input: finalInput,
276    }
277  
278    const canUseTool: CanUseToolFn = async () => ({
279      ...permissionResult,
280      decisionReason: {
281        type: 'mode',
282        mode: 'default' as const,
283      },
284    })
285  
286    // Add the assistant message with tool_use to messages BEFORE executing
287    // so the conversation history is complete (tool_use -> tool_result).
288    //
289    // On CCR resume, mutableMessages is seeded from the transcript and may already
290    // contain this tool_use. Pushing again would make normalizeMessagesForAPI merge
291    // same-ID assistants (concatenating content) and produce a duplicate tool_use
292    // ID, which the API rejects with "tool_use ids must be unique".
293    //
294    // Check for the specific tool_use_id rather than message.id: streaming yields
295    // each content block as a separate AssistantMessage sharing one message.id, so
296    // a [text, tool_use] response lands as two entries. filterUnresolvedToolUses may
297    // strip the tool_use entry but keep the text one; an id-based check would then
298    // wrongly skip the push while runTools below still executes, orphaning the result.
299    const alreadyPresent = mutableMessages.some(
300      m =>
301        m.type === 'assistant' &&
302        Array.isArray(m.message.content) &&
303        m.message.content.some(
304          b => b.type === 'tool_use' && 'id' in b && b.id === toolUseID,
305        ),
306    )
307    if (!alreadyPresent) {
308      mutableMessages.push(assistantMessage)
309      if (persistSession) {
310        await recordTranscript(mutableMessages)
311      }
312    }
313  
314    const sdkAssistantMessage: SDKMessage = {
315      ...assistantMessage,
316      session_id: getSessionId(),
317      parent_tool_use_id: null,
318    } as SDKMessage
319    yield sdkAssistantMessage
320  
321    // Execute the tool - errors are handled internally by runToolUse
322    for await (const update of runTools(
323      [finalToolUseBlock],
324      [assistantMessage],
325      canUseTool,
326      processUserInputContext,
327    )) {
328      if (update.message) {
329        mutableMessages.push(update.message)
330        if (persistSession) {
331          await recordTranscript(mutableMessages)
332        }
333  
334        const sdkMessage: SDKMessage = {
335          ...update.message,
336          session_id: getSessionId(),
337          parent_tool_use_id: null,
338        } as SDKMessage
339  
340        yield sdkMessage
341      }
342    }
343  }
344  
345  // Create a function to extract read files from messages
346  export function extractReadFilesFromMessages(
347    messages: Message[],
348    cwd: string,
349    maxSize: number = ASK_READ_FILE_STATE_CACHE_SIZE,
350  ): FileStateCache {
351    const cache = createFileStateCacheWithSizeLimit(maxSize)
352  
353    // First pass: find all FileReadTool/FileWriteTool/FileEditTool uses in assistant messages
354    const fileReadToolUseIds = new Map<string, string>() // toolUseId -> filePath
355    const fileWriteToolUseIds = new Map<
356      string,
357      { filePath: string; content: string }
358    >() // toolUseId -> { filePath, content }
359    const fileEditToolUseIds = new Map<string, string>() // toolUseId -> filePath
360  
361    for (const message of messages) {
362      if (
363        message.type === 'assistant' &&
364        Array.isArray(message.message.content)
365      ) {
366        for (const content of message.message.content) {
367          if (
368            content.type === 'tool_use' &&
369            content.name === FILE_READ_TOOL_NAME
370          ) {
371            // Extract file_path from the tool use input
372            const input = content.input as FileReadInput | undefined
373            // Ranged reads are not added to the cache.
374            if (
375              input?.file_path &&
376              input?.offset === undefined &&
377              input?.limit === undefined
378            ) {
379              // Normalize to absolute path for consistent cache lookups
380              const absolutePath = expandPath(input.file_path, cwd)
381              fileReadToolUseIds.set(content.id, absolutePath)
382            }
383          } else if (
384            content.type === 'tool_use' &&
385            content.name === FILE_WRITE_TOOL_NAME
386          ) {
387            // Extract file_path and content from the Write tool use input
388            const input = content.input as
389              | { file_path?: string; content?: string }
390              | undefined
391            if (input?.file_path && input?.content) {
392              // Normalize to absolute path for consistent cache lookups
393              const absolutePath = expandPath(input.file_path, cwd)
394              fileWriteToolUseIds.set(content.id, {
395                filePath: absolutePath,
396                content: input.content,
397              })
398            }
399          } else if (
400            content.type === 'tool_use' &&
401            content.name === FILE_EDIT_TOOL_NAME
402          ) {
403            // Edit's input has old_string/new_string, not the resulting content.
404            // Track the path so the second pass can read current disk state.
405            const input = content.input as { file_path?: string } | undefined
406            if (input?.file_path) {
407              const absolutePath = expandPath(input.file_path, cwd)
408              fileEditToolUseIds.set(content.id, absolutePath)
409            }
410          }
411        }
412      }
413    }
414  
415    // Second pass: find corresponding tool results and extract content
416    for (const message of messages) {
417      if (message.type === 'user' && Array.isArray(message.message.content)) {
418        for (const content of message.message.content) {
419          if (content.type === 'tool_result' && content.tool_use_id) {
420            // Handle Read tool results
421            const readFilePath = fileReadToolUseIds.get(content.tool_use_id)
422            if (
423              readFilePath &&
424              typeof content.content === 'string' &&
425              // Dedup stubs contain no file content — the earlier real Read
426              // already cached it. Chronological last-wins would otherwise
427              // overwrite the real entry with stub text.
428              !content.content.startsWith(FILE_UNCHANGED_STUB)
429            ) {
430              // Remove system-reminder blocks from the content
431              const processedContent = content.content.replace(
432                /<system-reminder>[\s\S]*?<\/system-reminder>/g,
433                '',
434              )
435  
436              // Extract the actual file content from the tool result
437              // Tool results for text files contain line numbers, we need to strip those
438              const fileContent = processedContent
439                .split('\n')
440                .map(stripLineNumberPrefix)
441                .join('\n')
442                .trim()
443  
444              // Cache the file content with the message timestamp
445              if (message.timestamp) {
446                const timestamp = new Date(message.timestamp).getTime()
447                cache.set(readFilePath, {
448                  content: fileContent,
449                  timestamp,
450                  offset: undefined,
451                  limit: undefined,
452                })
453              }
454            }
455  
456            // Handle Write tool results - use content from the tool input
457            const writeToolData = fileWriteToolUseIds.get(content.tool_use_id)
458            if (writeToolData && message.timestamp) {
459              const timestamp = new Date(message.timestamp).getTime()
460              cache.set(writeToolData.filePath, {
461                content: writeToolData.content,
462                timestamp,
463                offset: undefined,
464                limit: undefined,
465              })
466            }
467  
468            // Handle Edit tool results — post-edit content isn't in the
469            // tool_use input (only old_string/new_string) nor fully in the
470            // result (only a snippet). Read from disk now, using actual mtime
471            // so getChangedFiles's mtime check passes on the next turn.
472            //
473            // Callers seed the cache once at process start (print.ts --resume,
474            // Cowork cold-restart per turn), so disk content at extraction time
475            // IS the post-edit state. No dedup: processing every Edit preserves
476            // last-wins semantics when Read/Write interleave (Edit→Read→Edit).
477            const editFilePath = fileEditToolUseIds.get(content.tool_use_id)
478            if (editFilePath && content.is_error !== true) {
479              try {
480                const { content: diskContent } =
481                  readFileSyncWithMetadata(editFilePath)
482                cache.set(editFilePath, {
483                  content: diskContent,
484                  timestamp: getFileModificationTime(editFilePath),
485                  offset: undefined,
486                  limit: undefined,
487                })
488              } catch (e: unknown) {
489                if (!isFsInaccessible(e)) {
490                  throw e
491                }
492                // File deleted or inaccessible since the Edit — skip
493              }
494            }
495          }
496        }
497      }
498    }
499  
500    return cache
501  }
502  
503  /**
504   * Extract the top-level CLI tools used in BashTool calls from message history.
505   * Returns a deduplicated set of command names (e.g. 'vercel', 'aws', 'git').
506   */
507  export function extractBashToolsFromMessages(messages: Message[]): Set<string> {
508    const tools = new Set<string>()
509    for (const message of messages) {
510      if (
511        message.type === 'assistant' &&
512        Array.isArray(message.message.content)
513      ) {
514        for (const content of message.message.content) {
515          if (content.type === 'tool_use' && content.name === BASH_TOOL_NAME) {
516            const { input } = content
517            if (
518              typeof input !== 'object' ||
519              input === null ||
520              !('command' in input)
521            )
522              continue
523            const cmd = extractCliName(
524              typeof input.command === 'string' ? input.command : undefined,
525            )
526            if (cmd) {
527              tools.add(cmd)
528            }
529          }
530        }
531      }
532    }
533    return tools
534  }
535  
536  const STRIPPED_COMMANDS = new Set(['sudo'])
537  
538  /**
539   * Extract the actual CLI name from a bash command string, skipping
540   * env var assignments (e.g. `FOO=bar vercel` → `vercel`) and prefixes
541   * in STRIPPED_COMMANDS.
542   */
543  function extractCliName(command: string | undefined): string | undefined {
544    if (!command) return undefined
545    const tokens = command.trim().split(/\s+/)
546    for (const token of tokens) {
547      if (/^[A-Za-z_]\w*=/.test(token)) continue
548      if (STRIPPED_COMMANDS.has(token)) continue
549      return token
550    }
551    return undefined
552  }