/ QueryEngine.ts
QueryEngine.ts
1 import { feature } from 'bun:bundle' 2 import type { ContentBlockParam } from '@anthropic-ai/sdk/resources/messages.mjs' 3 import { randomUUID } from 'crypto' 4 import last from 'lodash-es/last.js' 5 import { 6 getSessionId, 7 isSessionPersistenceDisabled, 8 } from 'src/bootstrap/state.js' 9 import type { 10 PermissionMode, 11 SDKCompactBoundaryMessage, 12 SDKMessage, 13 SDKPermissionDenial, 14 SDKStatus, 15 SDKUserMessageReplay, 16 } from 'src/entrypoints/agentSdkTypes.js' 17 import { accumulateUsage, updateUsage } from 'src/services/api/claude.js' 18 import type { NonNullableUsage } from 'src/services/api/logging.js' 19 import { EMPTY_USAGE } from 'src/services/api/logging.js' 20 import stripAnsi from 'strip-ansi' 21 import type { Command } from './commands.js' 22 import { getSlashCommandToolSkills } from './commands.js' 23 import { 24 LOCAL_COMMAND_STDERR_TAG, 25 LOCAL_COMMAND_STDOUT_TAG, 26 } from './constants/xml.js' 27 import { 28 getModelUsage, 29 getTotalAPIDuration, 30 getTotalCost, 31 } from './cost-tracker.js' 32 import type { CanUseToolFn } from './hooks/useCanUseTool.js' 33 import { loadMemoryPrompt } from './memdir/memdir.js' 34 import { hasAutoMemPathOverride } from './memdir/paths.js' 35 import { query } from './query.js' 36 import { categorizeRetryableAPIError } from './services/api/errors.js' 37 import type { MCPServerConnection } from './services/mcp/types.js' 38 import type { AppState } from './state/AppState.js' 39 import { type Tools, type ToolUseContext, toolMatchesName } from './Tool.js' 40 import type { AgentDefinition } from './tools/AgentTool/loadAgentsDir.js' 41 import { SYNTHETIC_OUTPUT_TOOL_NAME } from './tools/SyntheticOutputTool/SyntheticOutputTool.js' 42 import type { Message } from './types/message.js' 43 import type { OrphanedPermission } from './types/textInputTypes.js' 44 import { createAbortController } from './utils/abortController.js' 45 import type { AttributionState } from './utils/commitAttribution.js' 46 import { getGlobalConfig } from './utils/config.js' 47 import { getCwd } from './utils/cwd.js' 48 import { isBareMode, isEnvTruthy } from './utils/envUtils.js' 49 import { getFastModeState } from './utils/fastMode.js' 50 import { 51 type FileHistoryState, 52 fileHistoryEnabled, 53 fileHistoryMakeSnapshot, 54 } from './utils/fileHistory.js' 55 import { 56 cloneFileStateCache, 57 type FileStateCache, 58 } from './utils/fileStateCache.js' 59 import { headlessProfilerCheckpoint } from './utils/headlessProfiler.js' 60 import { registerStructuredOutputEnforcement } from './utils/hooks/hookHelpers.js' 61 import { getInMemoryErrors } from './utils/log.js' 62 import { countToolCalls, SYNTHETIC_MESSAGES } from './utils/messages.js' 63 import { 64 getMainLoopModel, 65 parseUserSpecifiedModel, 66 } from './utils/model/model.js' 67 import { loadAllPluginsCacheOnly } from './utils/plugins/pluginLoader.js' 68 import { 69 type ProcessUserInputContext, 70 processUserInput, 71 } from './utils/processUserInput/processUserInput.js' 72 import { fetchSystemPromptParts } from './utils/queryContext.js' 73 import { setCwd } from './utils/Shell.js' 74 import { 75 flushSessionStorage, 76 recordTranscript, 77 } from './utils/sessionStorage.js' 78 import { asSystemPrompt } from './utils/systemPromptType.js' 79 import { resolveThemeSetting } from './utils/systemTheme.js' 80 import { 81 shouldEnableThinkingByDefault, 82 type ThinkingConfig, 83 } from './utils/thinking.js' 84 85 // Lazy: MessageSelector.tsx pulls React/ink; only needed for message filtering at query time 86 /* eslint-disable @typescript-eslint/no-require-imports */ 87 const messageSelector = 88 (): typeof import('src/components/MessageSelector.js') => 89 require('src/components/MessageSelector.js') 90 91 import { 92 localCommandOutputToSDKAssistantMessage, 93 toSDKCompactMetadata, 94 } from './utils/messages/mappers.js' 95 import { 96 buildSystemInitMessage, 97 sdkCompatToolName, 98 } from './utils/messages/systemInit.js' 99 import { 100 getScratchpadDir, 101 isScratchpadEnabled, 102 } from './utils/permissions/filesystem.js' 103 /* eslint-enable @typescript-eslint/no-require-imports */ 104 import { 105 handleOrphanedPermission, 106 isResultSuccessful, 107 normalizeMessage, 108 } from './utils/queryHelpers.js' 109 110 // Dead code elimination: conditional import for coordinator mode 111 /* eslint-disable @typescript-eslint/no-require-imports */ 112 const getCoordinatorUserContext: ( 113 mcpClients: ReadonlyArray<{ name: string }>, 114 scratchpadDir?: string, 115 ) => { [k: string]: string } = feature('COORDINATOR_MODE') 116 ? require('./coordinator/coordinatorMode.js').getCoordinatorUserContext 117 : () => ({}) 118 /* eslint-enable @typescript-eslint/no-require-imports */ 119 120 // Dead code elimination: conditional import for snip compaction 121 /* eslint-disable @typescript-eslint/no-require-imports */ 122 const snipModule = feature('HISTORY_SNIP') 123 ? (require('./services/compact/snipCompact.js') as typeof import('./services/compact/snipCompact.js')) 124 : null 125 const snipProjection = feature('HISTORY_SNIP') 126 ? (require('./services/compact/snipProjection.js') as typeof import('./services/compact/snipProjection.js')) 127 : null 128 /* eslint-enable @typescript-eslint/no-require-imports */ 129 130 export type QueryEngineConfig = { 131 cwd: string 132 tools: Tools 133 commands: Command[] 134 mcpClients: MCPServerConnection[] 135 agents: AgentDefinition[] 136 canUseTool: CanUseToolFn 137 getAppState: () => AppState 138 setAppState: (f: (prev: AppState) => AppState) => void 139 initialMessages?: Message[] 140 readFileCache: FileStateCache 141 customSystemPrompt?: string 142 appendSystemPrompt?: string 143 userSpecifiedModel?: string 144 fallbackModel?: string 145 thinkingConfig?: ThinkingConfig 146 maxTurns?: number 147 maxBudgetUsd?: number 148 taskBudget?: { total: number } 149 jsonSchema?: Record<string, unknown> 150 verbose?: boolean 151 replayUserMessages?: boolean 152 /** Handler for URL elicitations triggered by MCP tool -32042 errors. */ 153 handleElicitation?: ToolUseContext['handleElicitation'] 154 includePartialMessages?: boolean 155 setSDKStatus?: (status: SDKStatus) => void 156 abortController?: AbortController 157 orphanedPermission?: OrphanedPermission 158 /** 159 * Snip-boundary handler: receives each yielded system message plus the 160 * current mutableMessages store. Returns undefined if the message is not a 161 * snip boundary; otherwise returns the replayed snip result. Injected by 162 * ask() when HISTORY_SNIP is enabled so feature-gated strings stay inside 163 * the gated module (keeps QueryEngine free of excluded strings and testable 164 * despite feature() returning false under bun test). SDK-only: the REPL 165 * keeps full history for UI scrollback and projects on demand via 166 * projectSnippedView; QueryEngine truncates here to bound memory in long 167 * headless sessions (no UI to preserve). 168 */ 169 snipReplay?: ( 170 yieldedSystemMsg: Message, 171 store: Message[], 172 ) => { messages: Message[]; executed: boolean } | undefined 173 } 174 175 /** 176 * QueryEngine owns the query lifecycle and session state for a conversation. 177 * It extracts the core logic from ask() into a standalone class that can be 178 * used by both the headless/SDK path and (in a future phase) the REPL. 179 * 180 * One QueryEngine per conversation. Each submitMessage() call starts a new 181 * turn within the same conversation. State (messages, file cache, usage, etc.) 182 * persists across turns. 183 */ 184 export class QueryEngine { 185 private config: QueryEngineConfig 186 private mutableMessages: Message[] 187 private abortController: AbortController 188 private permissionDenials: SDKPermissionDenial[] 189 private totalUsage: NonNullableUsage 190 private hasHandledOrphanedPermission = false 191 private readFileState: FileStateCache 192 // Turn-scoped skill discovery tracking (feeds was_discovered on 193 // tengu_skill_tool_invocation). Must persist across the two 194 // processUserInputContext rebuilds inside submitMessage, but is cleared 195 // at the start of each submitMessage to avoid unbounded growth across 196 // many turns in SDK mode. 197 private discoveredSkillNames = new Set<string>() 198 private loadedNestedMemoryPaths = new Set<string>() 199 200 constructor(config: QueryEngineConfig) { 201 this.config = config 202 this.mutableMessages = config.initialMessages ?? [] 203 this.abortController = config.abortController ?? createAbortController() 204 this.permissionDenials = [] 205 this.readFileState = config.readFileCache 206 this.totalUsage = EMPTY_USAGE 207 } 208 209 async *submitMessage( 210 prompt: string | ContentBlockParam[], 211 options?: { uuid?: string; isMeta?: boolean }, 212 ): AsyncGenerator<SDKMessage, void, unknown> { 213 const { 214 cwd, 215 commands, 216 tools, 217 mcpClients, 218 verbose = false, 219 thinkingConfig, 220 maxTurns, 221 maxBudgetUsd, 222 taskBudget, 223 canUseTool, 224 customSystemPrompt, 225 appendSystemPrompt, 226 userSpecifiedModel, 227 fallbackModel, 228 jsonSchema, 229 getAppState, 230 setAppState, 231 replayUserMessages = false, 232 includePartialMessages = false, 233 agents = [], 234 setSDKStatus, 235 orphanedPermission, 236 } = this.config 237 238 this.discoveredSkillNames.clear() 239 setCwd(cwd) 240 const persistSession = !isSessionPersistenceDisabled() 241 const startTime = Date.now() 242 243 // Wrap canUseTool to track permission denials 244 const wrappedCanUseTool: CanUseToolFn = async ( 245 tool, 246 input, 247 toolUseContext, 248 assistantMessage, 249 toolUseID, 250 forceDecision, 251 ) => { 252 const result = await canUseTool( 253 tool, 254 input, 255 toolUseContext, 256 assistantMessage, 257 toolUseID, 258 forceDecision, 259 ) 260 261 // Track denials for SDK reporting 262 if (result.behavior !== 'allow') { 263 this.permissionDenials.push({ 264 tool_name: sdkCompatToolName(tool.name), 265 tool_use_id: toolUseID, 266 tool_input: input, 267 }) 268 } 269 270 return result 271 } 272 273 const initialAppState = getAppState() 274 const initialMainLoopModel = userSpecifiedModel 275 ? parseUserSpecifiedModel(userSpecifiedModel) 276 : getMainLoopModel() 277 278 const initialThinkingConfig: ThinkingConfig = thinkingConfig 279 ? thinkingConfig 280 : shouldEnableThinkingByDefault() !== false 281 ? { type: 'adaptive' } 282 : { type: 'disabled' } 283 284 headlessProfilerCheckpoint('before_getSystemPrompt') 285 // Narrow once so TS tracks the type through the conditionals below. 286 const customPrompt = 287 typeof customSystemPrompt === 'string' ? customSystemPrompt : undefined 288 const { 289 defaultSystemPrompt, 290 userContext: baseUserContext, 291 systemContext, 292 } = await fetchSystemPromptParts({ 293 tools, 294 mainLoopModel: initialMainLoopModel, 295 additionalWorkingDirectories: Array.from( 296 initialAppState.toolPermissionContext.additionalWorkingDirectories.keys(), 297 ), 298 mcpClients, 299 customSystemPrompt: customPrompt, 300 }) 301 headlessProfilerCheckpoint('after_getSystemPrompt') 302 const userContext = { 303 ...baseUserContext, 304 ...getCoordinatorUserContext( 305 mcpClients, 306 isScratchpadEnabled() ? getScratchpadDir() : undefined, 307 ), 308 } 309 310 // When an SDK caller provides a custom system prompt AND has set 311 // CLAUDE_COWORK_MEMORY_PATH_OVERRIDE, inject the memory-mechanics prompt. 312 // The env var is an explicit opt-in signal — the caller has wired up 313 // a memory directory and needs Claude to know how to use it (which 314 // Write/Edit tools to call, MEMORY.md filename, loading semantics). 315 // The caller can layer their own policy text via appendSystemPrompt. 316 const memoryMechanicsPrompt = 317 customPrompt !== undefined && hasAutoMemPathOverride() 318 ? await loadMemoryPrompt() 319 : null 320 321 const systemPrompt = asSystemPrompt([ 322 ...(customPrompt !== undefined ? [customPrompt] : defaultSystemPrompt), 323 ...(memoryMechanicsPrompt ? [memoryMechanicsPrompt] : []), 324 ...(appendSystemPrompt ? [appendSystemPrompt] : []), 325 ]) 326 327 // Register function hook for structured output enforcement 328 const hasStructuredOutputTool = tools.some(t => 329 toolMatchesName(t, SYNTHETIC_OUTPUT_TOOL_NAME), 330 ) 331 if (jsonSchema && hasStructuredOutputTool) { 332 registerStructuredOutputEnforcement(setAppState, getSessionId()) 333 } 334 335 let processUserInputContext: ProcessUserInputContext = { 336 messages: this.mutableMessages, 337 // Slash commands that mutate the message array (e.g. /force-snip) 338 // call setMessages(fn). In interactive mode this writes back to 339 // AppState; in print mode we write back to mutableMessages so the 340 // rest of the query loop (push at :389, snapshot at :392) sees 341 // the result. The second processUserInputContext below (after 342 // slash-command processing) keeps the no-op — nothing else calls 343 // setMessages past that point. 344 setMessages: fn => { 345 this.mutableMessages = fn(this.mutableMessages) 346 }, 347 onChangeAPIKey: () => {}, 348 handleElicitation: this.config.handleElicitation, 349 options: { 350 commands, 351 debug: false, // we use stdout, so don't want to clobber it 352 tools, 353 verbose, 354 mainLoopModel: initialMainLoopModel, 355 thinkingConfig: initialThinkingConfig, 356 mcpClients, 357 mcpResources: {}, 358 ideInstallationStatus: null, 359 isNonInteractiveSession: true, 360 customSystemPrompt, 361 appendSystemPrompt, 362 agentDefinitions: { activeAgents: agents, allAgents: [] }, 363 theme: resolveThemeSetting(getGlobalConfig().theme), 364 maxBudgetUsd, 365 }, 366 getAppState, 367 setAppState, 368 abortController: this.abortController, 369 readFileState: this.readFileState, 370 nestedMemoryAttachmentTriggers: new Set<string>(), 371 loadedNestedMemoryPaths: this.loadedNestedMemoryPaths, 372 dynamicSkillDirTriggers: new Set<string>(), 373 discoveredSkillNames: this.discoveredSkillNames, 374 setInProgressToolUseIDs: () => {}, 375 setResponseLength: () => {}, 376 updateFileHistoryState: ( 377 updater: (prev: FileHistoryState) => FileHistoryState, 378 ) => { 379 setAppState(prev => { 380 const updated = updater(prev.fileHistory) 381 if (updated === prev.fileHistory) return prev 382 return { ...prev, fileHistory: updated } 383 }) 384 }, 385 updateAttributionState: ( 386 updater: (prev: AttributionState) => AttributionState, 387 ) => { 388 setAppState(prev => { 389 const updated = updater(prev.attribution) 390 if (updated === prev.attribution) return prev 391 return { ...prev, attribution: updated } 392 }) 393 }, 394 setSDKStatus, 395 } 396 397 // Handle orphaned permission (only once per engine lifetime) 398 if (orphanedPermission && !this.hasHandledOrphanedPermission) { 399 this.hasHandledOrphanedPermission = true 400 for await (const message of handleOrphanedPermission( 401 orphanedPermission, 402 tools, 403 this.mutableMessages, 404 processUserInputContext, 405 )) { 406 yield message 407 } 408 } 409 410 const { 411 messages: messagesFromUserInput, 412 shouldQuery, 413 allowedTools, 414 model: modelFromUserInput, 415 resultText, 416 } = await processUserInput({ 417 input: prompt, 418 mode: 'prompt', 419 setToolJSX: () => {}, 420 context: { 421 ...processUserInputContext, 422 messages: this.mutableMessages, 423 }, 424 messages: this.mutableMessages, 425 uuid: options?.uuid, 426 isMeta: options?.isMeta, 427 querySource: 'sdk', 428 }) 429 430 // Push new messages, including user input and any attachments 431 this.mutableMessages.push(...messagesFromUserInput) 432 433 // Update params to reflect updates from processing /slash commands 434 const messages = [...this.mutableMessages] 435 436 // Persist the user's message(s) to transcript BEFORE entering the query 437 // loop. The for-await below only calls recordTranscript when ask() yields 438 // an assistant/user/compact_boundary message — which doesn't happen until 439 // the API responds. If the process is killed before that (e.g. user clicks 440 // Stop in cowork seconds after send), the transcript is left with only 441 // queue-operation entries; getLastSessionLog filters those out, returns 442 // null, and --resume fails with "No conversation found". Writing now makes 443 // the transcript resumable from the point the user message was accepted, 444 // even if no API response ever arrives. 445 // 446 // --bare / SIMPLE: fire-and-forget. Scripted calls don't --resume after 447 // kill-mid-request. The await is ~4ms on SSD, ~30ms under disk contention 448 // — the single largest controllable critical-path cost after module eval. 449 // Transcript is still written (for post-hoc debugging); just not blocking. 450 if (persistSession && messagesFromUserInput.length > 0) { 451 const transcriptPromise = recordTranscript(messages) 452 if (isBareMode()) { 453 void transcriptPromise 454 } else { 455 await transcriptPromise 456 if ( 457 isEnvTruthy(process.env.CLAUDE_CODE_EAGER_FLUSH) || 458 isEnvTruthy(process.env.CLAUDE_CODE_IS_COWORK) 459 ) { 460 await flushSessionStorage() 461 } 462 } 463 } 464 465 // Filter messages that should be acknowledged after transcript 466 const replayableMessages = messagesFromUserInput.filter( 467 msg => 468 (msg.type === 'user' && 469 !msg.isMeta && // Skip synthetic caveat messages 470 !msg.toolUseResult && // Skip tool results (they'll be acked from query) 471 messageSelector().selectableUserMessagesFilter(msg)) || // Skip non-user-authored messages (task notifications, etc.) 472 (msg.type === 'system' && msg.subtype === 'compact_boundary'), // Always ack compact boundaries 473 ) 474 const messagesToAck = replayUserMessages ? replayableMessages : [] 475 476 // Update the ToolPermissionContext based on user input processing (as necessary) 477 setAppState(prev => ({ 478 ...prev, 479 toolPermissionContext: { 480 ...prev.toolPermissionContext, 481 alwaysAllowRules: { 482 ...prev.toolPermissionContext.alwaysAllowRules, 483 command: allowedTools, 484 }, 485 }, 486 })) 487 488 const mainLoopModel = modelFromUserInput ?? initialMainLoopModel 489 490 // Recreate after processing the prompt to pick up updated messages and 491 // model (from slash commands). 492 processUserInputContext = { 493 messages, 494 setMessages: () => {}, 495 onChangeAPIKey: () => {}, 496 handleElicitation: this.config.handleElicitation, 497 options: { 498 commands, 499 debug: false, 500 tools, 501 verbose, 502 mainLoopModel, 503 thinkingConfig: initialThinkingConfig, 504 mcpClients, 505 mcpResources: {}, 506 ideInstallationStatus: null, 507 isNonInteractiveSession: true, 508 customSystemPrompt, 509 appendSystemPrompt, 510 theme: resolveThemeSetting(getGlobalConfig().theme), 511 agentDefinitions: { activeAgents: agents, allAgents: [] }, 512 maxBudgetUsd, 513 }, 514 getAppState, 515 setAppState, 516 abortController: this.abortController, 517 readFileState: this.readFileState, 518 nestedMemoryAttachmentTriggers: new Set<string>(), 519 loadedNestedMemoryPaths: this.loadedNestedMemoryPaths, 520 dynamicSkillDirTriggers: new Set<string>(), 521 discoveredSkillNames: this.discoveredSkillNames, 522 setInProgressToolUseIDs: () => {}, 523 setResponseLength: () => {}, 524 updateFileHistoryState: processUserInputContext.updateFileHistoryState, 525 updateAttributionState: processUserInputContext.updateAttributionState, 526 setSDKStatus, 527 } 528 529 headlessProfilerCheckpoint('before_skills_plugins') 530 // Cache-only: headless/SDK/CCR startup must not block on network for 531 // ref-tracked plugins. CCR populates the cache via CLAUDE_CODE_SYNC_PLUGIN_INSTALL 532 // (headlessPluginInstall) or CLAUDE_CODE_PLUGIN_SEED_DIR before this runs; 533 // SDK callers that need fresh source can call /reload-plugins. 534 const [skills, { enabled: enabledPlugins }] = await Promise.all([ 535 getSlashCommandToolSkills(getCwd()), 536 loadAllPluginsCacheOnly(), 537 ]) 538 headlessProfilerCheckpoint('after_skills_plugins') 539 540 yield buildSystemInitMessage({ 541 tools, 542 mcpClients, 543 model: mainLoopModel, 544 permissionMode: initialAppState.toolPermissionContext 545 .mode as PermissionMode, // TODO: avoid the cast 546 commands, 547 agents, 548 skills, 549 plugins: enabledPlugins, 550 fastMode: initialAppState.fastMode, 551 }) 552 553 // Record when system message is yielded for headless latency tracking 554 headlessProfilerCheckpoint('system_message_yielded') 555 556 if (!shouldQuery) { 557 // Return the results of local slash commands. 558 // Use messagesFromUserInput (not replayableMessages) for command output 559 // because selectableUserMessagesFilter excludes local-command-stdout tags. 560 for (const msg of messagesFromUserInput) { 561 if ( 562 msg.type === 'user' && 563 typeof msg.message.content === 'string' && 564 (msg.message.content.includes(`<${LOCAL_COMMAND_STDOUT_TAG}>`) || 565 msg.message.content.includes(`<${LOCAL_COMMAND_STDERR_TAG}>`) || 566 msg.isCompactSummary) 567 ) { 568 yield { 569 type: 'user', 570 message: { 571 ...msg.message, 572 content: stripAnsi(msg.message.content), 573 }, 574 session_id: getSessionId(), 575 parent_tool_use_id: null, 576 uuid: msg.uuid, 577 timestamp: msg.timestamp, 578 isReplay: !msg.isCompactSummary, 579 isSynthetic: msg.isMeta || msg.isVisibleInTranscriptOnly, 580 } as SDKUserMessageReplay 581 } 582 583 // Local command output — yield as a synthetic assistant message so 584 // RC renders it as assistant-style text rather than a user bubble. 585 // Emitted as assistant (not the dedicated SDKLocalCommandOutputMessage 586 // system subtype) so mobile clients + session-ingress can parse it. 587 if ( 588 msg.type === 'system' && 589 msg.subtype === 'local_command' && 590 typeof msg.content === 'string' && 591 (msg.content.includes(`<${LOCAL_COMMAND_STDOUT_TAG}>`) || 592 msg.content.includes(`<${LOCAL_COMMAND_STDERR_TAG}>`)) 593 ) { 594 yield localCommandOutputToSDKAssistantMessage(msg.content, msg.uuid) 595 } 596 597 if (msg.type === 'system' && msg.subtype === 'compact_boundary') { 598 yield { 599 type: 'system', 600 subtype: 'compact_boundary' as const, 601 session_id: getSessionId(), 602 uuid: msg.uuid, 603 compact_metadata: toSDKCompactMetadata(msg.compactMetadata), 604 } as SDKCompactBoundaryMessage 605 } 606 } 607 608 if (persistSession) { 609 await recordTranscript(messages) 610 if ( 611 isEnvTruthy(process.env.CLAUDE_CODE_EAGER_FLUSH) || 612 isEnvTruthy(process.env.CLAUDE_CODE_IS_COWORK) 613 ) { 614 await flushSessionStorage() 615 } 616 } 617 618 yield { 619 type: 'result', 620 subtype: 'success', 621 is_error: false, 622 duration_ms: Date.now() - startTime, 623 duration_api_ms: getTotalAPIDuration(), 624 num_turns: messages.length - 1, 625 result: resultText ?? '', 626 stop_reason: null, 627 session_id: getSessionId(), 628 total_cost_usd: getTotalCost(), 629 usage: this.totalUsage, 630 modelUsage: getModelUsage(), 631 permission_denials: this.permissionDenials, 632 fast_mode_state: getFastModeState( 633 mainLoopModel, 634 initialAppState.fastMode, 635 ), 636 uuid: randomUUID(), 637 } 638 return 639 } 640 641 if (fileHistoryEnabled() && persistSession) { 642 messagesFromUserInput 643 .filter(messageSelector().selectableUserMessagesFilter) 644 .forEach(message => { 645 void fileHistoryMakeSnapshot( 646 (updater: (prev: FileHistoryState) => FileHistoryState) => { 647 setAppState(prev => ({ 648 ...prev, 649 fileHistory: updater(prev.fileHistory), 650 })) 651 }, 652 message.uuid, 653 ) 654 }) 655 } 656 657 // Track current message usage (reset on each message_start) 658 let currentMessageUsage: NonNullableUsage = EMPTY_USAGE 659 let turnCount = 1 660 let hasAcknowledgedInitialMessages = false 661 // Track structured output from StructuredOutput tool calls 662 let structuredOutputFromTool: unknown 663 // Track the last stop_reason from assistant messages 664 let lastStopReason: string | null = null 665 // Reference-based watermark so error_during_execution's errors[] is 666 // turn-scoped. A length-based index breaks when the 100-entry ring buffer 667 // shift()s during the turn — the index slides. If this entry is rotated 668 // out, lastIndexOf returns -1 and we include everything (safe fallback). 669 const errorLogWatermark = getInMemoryErrors().at(-1) 670 // Snapshot count before this query for delta-based retry limiting 671 const initialStructuredOutputCalls = jsonSchema 672 ? countToolCalls(this.mutableMessages, SYNTHETIC_OUTPUT_TOOL_NAME) 673 : 0 674 675 for await (const message of query({ 676 messages, 677 systemPrompt, 678 userContext, 679 systemContext, 680 canUseTool: wrappedCanUseTool, 681 toolUseContext: processUserInputContext, 682 fallbackModel, 683 querySource: 'sdk', 684 maxTurns, 685 taskBudget, 686 })) { 687 // Record assistant, user, and compact boundary messages 688 if ( 689 message.type === 'assistant' || 690 message.type === 'user' || 691 (message.type === 'system' && message.subtype === 'compact_boundary') 692 ) { 693 // Before writing a compact boundary, flush any in-memory-only 694 // messages up through the preservedSegment tail. Attachments and 695 // progress are now recorded inline (their switch cases below), but 696 // this flush still matters for the preservedSegment tail walk. 697 // If the SDK subprocess restarts before then (claude-desktop kills 698 // between turns), tailUuid points to a never-written message → 699 // applyPreservedSegmentRelinks fails its tail→head walk → returns 700 // without pruning → resume loads full pre-compact history. 701 if ( 702 persistSession && 703 message.type === 'system' && 704 message.subtype === 'compact_boundary' 705 ) { 706 const tailUuid = message.compactMetadata?.preservedSegment?.tailUuid 707 if (tailUuid) { 708 const tailIdx = this.mutableMessages.findLastIndex( 709 m => m.uuid === tailUuid, 710 ) 711 if (tailIdx !== -1) { 712 await recordTranscript(this.mutableMessages.slice(0, tailIdx + 1)) 713 } 714 } 715 } 716 messages.push(message) 717 if (persistSession) { 718 // Fire-and-forget for assistant messages. claude.ts yields one 719 // assistant message per content block, then mutates the last 720 // one's message.usage/stop_reason on message_delta — relying on 721 // the write queue's 100ms lazy jsonStringify. Awaiting here 722 // blocks ask()'s generator, so message_delta can't run until 723 // every block is consumed; the drain timer (started at block 1) 724 // elapses first. Interactive CC doesn't hit this because 725 // useLogMessages.ts fire-and-forgets. enqueueWrite is 726 // order-preserving so fire-and-forget here is safe. 727 if (message.type === 'assistant') { 728 void recordTranscript(messages) 729 } else { 730 await recordTranscript(messages) 731 } 732 } 733 734 // Acknowledge initial user messages after first transcript recording 735 if (!hasAcknowledgedInitialMessages && messagesToAck.length > 0) { 736 hasAcknowledgedInitialMessages = true 737 for (const msgToAck of messagesToAck) { 738 if (msgToAck.type === 'user') { 739 yield { 740 type: 'user', 741 message: msgToAck.message, 742 session_id: getSessionId(), 743 parent_tool_use_id: null, 744 uuid: msgToAck.uuid, 745 timestamp: msgToAck.timestamp, 746 isReplay: true, 747 } as SDKUserMessageReplay 748 } 749 } 750 } 751 } 752 753 if (message.type === 'user') { 754 turnCount++ 755 } 756 757 switch (message.type) { 758 case 'tombstone': 759 // Tombstone messages are control signals for removing messages, skip them 760 break 761 case 'assistant': 762 // Capture stop_reason if already set (synthetic messages). For 763 // streamed responses, this is null at content_block_stop time; 764 // the real value arrives via message_delta (handled below). 765 if (message.message.stop_reason != null) { 766 lastStopReason = message.message.stop_reason 767 } 768 this.mutableMessages.push(message) 769 yield* normalizeMessage(message) 770 break 771 case 'progress': 772 this.mutableMessages.push(message) 773 // Record inline so the dedup loop in the next ask() call sees it 774 // as already-recorded. Without this, deferred progress interleaves 775 // with already-recorded tool_results in mutableMessages, and the 776 // dedup walk freezes startingParentUuid at the wrong message — 777 // forking the chain and orphaning the conversation on resume. 778 if (persistSession) { 779 messages.push(message) 780 void recordTranscript(messages) 781 } 782 yield* normalizeMessage(message) 783 break 784 case 'user': 785 this.mutableMessages.push(message) 786 yield* normalizeMessage(message) 787 break 788 case 'stream_event': 789 if (message.event.type === 'message_start') { 790 // Reset current message usage for new message 791 currentMessageUsage = EMPTY_USAGE 792 currentMessageUsage = updateUsage( 793 currentMessageUsage, 794 message.event.message.usage, 795 ) 796 } 797 if (message.event.type === 'message_delta') { 798 currentMessageUsage = updateUsage( 799 currentMessageUsage, 800 message.event.usage, 801 ) 802 // Capture stop_reason from message_delta. The assistant message 803 // is yielded at content_block_stop with stop_reason=null; the 804 // real value only arrives here (see claude.ts message_delta 805 // handler). Without this, result.stop_reason is always null. 806 if (message.event.delta.stop_reason != null) { 807 lastStopReason = message.event.delta.stop_reason 808 } 809 } 810 if (message.event.type === 'message_stop') { 811 // Accumulate current message usage into total 812 this.totalUsage = accumulateUsage( 813 this.totalUsage, 814 currentMessageUsage, 815 ) 816 } 817 818 if (includePartialMessages) { 819 yield { 820 type: 'stream_event' as const, 821 event: message.event, 822 session_id: getSessionId(), 823 parent_tool_use_id: null, 824 uuid: randomUUID(), 825 } 826 } 827 828 break 829 case 'attachment': 830 this.mutableMessages.push(message) 831 // Record inline (same reason as progress above). 832 if (persistSession) { 833 messages.push(message) 834 void recordTranscript(messages) 835 } 836 837 // Extract structured output from StructuredOutput tool calls 838 if (message.attachment.type === 'structured_output') { 839 structuredOutputFromTool = message.attachment.data 840 } 841 // Handle max turns reached signal from query.ts 842 else if (message.attachment.type === 'max_turns_reached') { 843 if (persistSession) { 844 if ( 845 isEnvTruthy(process.env.CLAUDE_CODE_EAGER_FLUSH) || 846 isEnvTruthy(process.env.CLAUDE_CODE_IS_COWORK) 847 ) { 848 await flushSessionStorage() 849 } 850 } 851 yield { 852 type: 'result', 853 subtype: 'error_max_turns', 854 duration_ms: Date.now() - startTime, 855 duration_api_ms: getTotalAPIDuration(), 856 is_error: true, 857 num_turns: message.attachment.turnCount, 858 stop_reason: lastStopReason, 859 session_id: getSessionId(), 860 total_cost_usd: getTotalCost(), 861 usage: this.totalUsage, 862 modelUsage: getModelUsage(), 863 permission_denials: this.permissionDenials, 864 fast_mode_state: getFastModeState( 865 mainLoopModel, 866 initialAppState.fastMode, 867 ), 868 uuid: randomUUID(), 869 errors: [ 870 `Reached maximum number of turns (${message.attachment.maxTurns})`, 871 ], 872 } 873 return 874 } 875 // Yield queued_command attachments as SDK user message replays 876 else if ( 877 replayUserMessages && 878 message.attachment.type === 'queued_command' 879 ) { 880 yield { 881 type: 'user', 882 message: { 883 role: 'user' as const, 884 content: message.attachment.prompt, 885 }, 886 session_id: getSessionId(), 887 parent_tool_use_id: null, 888 uuid: message.attachment.source_uuid || message.uuid, 889 timestamp: message.timestamp, 890 isReplay: true, 891 } as SDKUserMessageReplay 892 } 893 break 894 case 'stream_request_start': 895 // Don't yield stream request start messages 896 break 897 case 'system': { 898 // Snip boundary: replay on our store to remove zombie messages and 899 // stale markers. The yielded boundary is a signal, not data to push — 900 // the replay produces its own equivalent boundary. Without this, 901 // markers persist and re-trigger on every turn, and mutableMessages 902 // never shrinks (memory leak in long SDK sessions). The subtype 903 // check lives inside the injected callback so feature-gated strings 904 // stay out of this file (excluded-strings check). 905 const snipResult = this.config.snipReplay?.( 906 message, 907 this.mutableMessages, 908 ) 909 if (snipResult !== undefined) { 910 if (snipResult.executed) { 911 this.mutableMessages.length = 0 912 this.mutableMessages.push(...snipResult.messages) 913 } 914 break 915 } 916 this.mutableMessages.push(message) 917 // Yield compact boundary messages to SDK 918 if ( 919 message.subtype === 'compact_boundary' && 920 message.compactMetadata 921 ) { 922 // Release pre-compaction messages for GC. The boundary was just 923 // pushed so it's the last element. query.ts already uses 924 // getMessagesAfterCompactBoundary() internally, so only 925 // post-boundary messages are needed going forward. 926 const mutableBoundaryIdx = this.mutableMessages.length - 1 927 if (mutableBoundaryIdx > 0) { 928 this.mutableMessages.splice(0, mutableBoundaryIdx) 929 } 930 const localBoundaryIdx = messages.length - 1 931 if (localBoundaryIdx > 0) { 932 messages.splice(0, localBoundaryIdx) 933 } 934 935 yield { 936 type: 'system', 937 subtype: 'compact_boundary' as const, 938 session_id: getSessionId(), 939 uuid: message.uuid, 940 compact_metadata: toSDKCompactMetadata(message.compactMetadata), 941 } 942 } 943 if (message.subtype === 'api_error') { 944 yield { 945 type: 'system', 946 subtype: 'api_retry' as const, 947 attempt: message.retryAttempt, 948 max_retries: message.maxRetries, 949 retry_delay_ms: message.retryInMs, 950 error_status: message.error.status ?? null, 951 error: categorizeRetryableAPIError(message.error), 952 session_id: getSessionId(), 953 uuid: message.uuid, 954 } 955 } 956 // Don't yield other system messages in headless mode 957 break 958 } 959 case 'tool_use_summary': 960 // Yield tool use summary messages to SDK 961 yield { 962 type: 'tool_use_summary' as const, 963 summary: message.summary, 964 preceding_tool_use_ids: message.precedingToolUseIds, 965 session_id: getSessionId(), 966 uuid: message.uuid, 967 } 968 break 969 } 970 971 // Check if USD budget has been exceeded 972 if (maxBudgetUsd !== undefined && getTotalCost() >= maxBudgetUsd) { 973 if (persistSession) { 974 if ( 975 isEnvTruthy(process.env.CLAUDE_CODE_EAGER_FLUSH) || 976 isEnvTruthy(process.env.CLAUDE_CODE_IS_COWORK) 977 ) { 978 await flushSessionStorage() 979 } 980 } 981 yield { 982 type: 'result', 983 subtype: 'error_max_budget_usd', 984 duration_ms: Date.now() - startTime, 985 duration_api_ms: getTotalAPIDuration(), 986 is_error: true, 987 num_turns: turnCount, 988 stop_reason: lastStopReason, 989 session_id: getSessionId(), 990 total_cost_usd: getTotalCost(), 991 usage: this.totalUsage, 992 modelUsage: getModelUsage(), 993 permission_denials: this.permissionDenials, 994 fast_mode_state: getFastModeState( 995 mainLoopModel, 996 initialAppState.fastMode, 997 ), 998 uuid: randomUUID(), 999 errors: [`Reached maximum budget ($${maxBudgetUsd})`], 1000 } 1001 return 1002 } 1003 1004 // Check if structured output retry limit exceeded (only on user messages) 1005 if (message.type === 'user' && jsonSchema) { 1006 const currentCalls = countToolCalls( 1007 this.mutableMessages, 1008 SYNTHETIC_OUTPUT_TOOL_NAME, 1009 ) 1010 const callsThisQuery = currentCalls - initialStructuredOutputCalls 1011 const maxRetries = parseInt( 1012 process.env.MAX_STRUCTURED_OUTPUT_RETRIES || '5', 1013 10, 1014 ) 1015 if (callsThisQuery >= maxRetries) { 1016 if (persistSession) { 1017 if ( 1018 isEnvTruthy(process.env.CLAUDE_CODE_EAGER_FLUSH) || 1019 isEnvTruthy(process.env.CLAUDE_CODE_IS_COWORK) 1020 ) { 1021 await flushSessionStorage() 1022 } 1023 } 1024 yield { 1025 type: 'result', 1026 subtype: 'error_max_structured_output_retries', 1027 duration_ms: Date.now() - startTime, 1028 duration_api_ms: getTotalAPIDuration(), 1029 is_error: true, 1030 num_turns: turnCount, 1031 stop_reason: lastStopReason, 1032 session_id: getSessionId(), 1033 total_cost_usd: getTotalCost(), 1034 usage: this.totalUsage, 1035 modelUsage: getModelUsage(), 1036 permission_denials: this.permissionDenials, 1037 fast_mode_state: getFastModeState( 1038 mainLoopModel, 1039 initialAppState.fastMode, 1040 ), 1041 uuid: randomUUID(), 1042 errors: [ 1043 `Failed to provide valid structured output after ${maxRetries} attempts`, 1044 ], 1045 } 1046 return 1047 } 1048 } 1049 } 1050 1051 // Stop hooks yield progress/attachment messages AFTER the assistant 1052 // response (via yield* handleStopHooks in query.ts). Since #23537 pushes 1053 // those to `messages` inline, last(messages) can be a progress/attachment 1054 // instead of the assistant — which makes textResult extraction below 1055 // return '' and -p mode emit a blank line. Allowlist to assistant|user: 1056 // isResultSuccessful handles both (user with all tool_result blocks is a 1057 // valid successful terminal state). 1058 const result = messages.findLast( 1059 m => m.type === 'assistant' || m.type === 'user', 1060 ) 1061 // Capture for the error_during_execution diagnostic — isResultSuccessful 1062 // is a type predicate (message is Message), so inside the false branch 1063 // `result` narrows to never and these accesses don't typecheck. 1064 const edeResultType = result?.type ?? 'undefined' 1065 const edeLastContentType = 1066 result?.type === 'assistant' 1067 ? (last(result.message.content)?.type ?? 'none') 1068 : 'n/a' 1069 1070 // Flush buffered transcript writes before yielding result. 1071 // The desktop app kills the CLI process immediately after receiving the 1072 // result message, so any unflushed writes would be lost. 1073 if (persistSession) { 1074 if ( 1075 isEnvTruthy(process.env.CLAUDE_CODE_EAGER_FLUSH) || 1076 isEnvTruthy(process.env.CLAUDE_CODE_IS_COWORK) 1077 ) { 1078 await flushSessionStorage() 1079 } 1080 } 1081 1082 if (!isResultSuccessful(result, lastStopReason)) { 1083 yield { 1084 type: 'result', 1085 subtype: 'error_during_execution', 1086 duration_ms: Date.now() - startTime, 1087 duration_api_ms: getTotalAPIDuration(), 1088 is_error: true, 1089 num_turns: turnCount, 1090 stop_reason: lastStopReason, 1091 session_id: getSessionId(), 1092 total_cost_usd: getTotalCost(), 1093 usage: this.totalUsage, 1094 modelUsage: getModelUsage(), 1095 permission_denials: this.permissionDenials, 1096 fast_mode_state: getFastModeState( 1097 mainLoopModel, 1098 initialAppState.fastMode, 1099 ), 1100 uuid: randomUUID(), 1101 // Diagnostic prefix: these are what isResultSuccessful() checks — if 1102 // the result type isn't assistant-with-text/thinking or user-with- 1103 // tool_result, and stop_reason isn't end_turn, that's why this fired. 1104 // errors[] is turn-scoped via the watermark; previously it dumped the 1105 // entire process's logError buffer (ripgrep timeouts, ENOENT, etc). 1106 errors: (() => { 1107 const all = getInMemoryErrors() 1108 const start = errorLogWatermark 1109 ? all.lastIndexOf(errorLogWatermark) + 1 1110 : 0 1111 return [ 1112 `[ede_diagnostic] result_type=${edeResultType} last_content_type=${edeLastContentType} stop_reason=${lastStopReason}`, 1113 ...all.slice(start).map(_ => _.error), 1114 ] 1115 })(), 1116 } 1117 return 1118 } 1119 1120 // Extract the text result based on message type 1121 let textResult = '' 1122 let isApiError = false 1123 1124 if (result.type === 'assistant') { 1125 const lastContent = last(result.message.content) 1126 if ( 1127 lastContent?.type === 'text' && 1128 !SYNTHETIC_MESSAGES.has(lastContent.text) 1129 ) { 1130 textResult = lastContent.text 1131 } 1132 isApiError = Boolean(result.isApiErrorMessage) 1133 } 1134 1135 yield { 1136 type: 'result', 1137 subtype: 'success', 1138 is_error: isApiError, 1139 duration_ms: Date.now() - startTime, 1140 duration_api_ms: getTotalAPIDuration(), 1141 num_turns: turnCount, 1142 result: textResult, 1143 stop_reason: lastStopReason, 1144 session_id: getSessionId(), 1145 total_cost_usd: getTotalCost(), 1146 usage: this.totalUsage, 1147 modelUsage: getModelUsage(), 1148 permission_denials: this.permissionDenials, 1149 structured_output: structuredOutputFromTool, 1150 fast_mode_state: getFastModeState( 1151 mainLoopModel, 1152 initialAppState.fastMode, 1153 ), 1154 uuid: randomUUID(), 1155 } 1156 } 1157 1158 interrupt(): void { 1159 this.abortController.abort() 1160 } 1161 1162 getMessages(): readonly Message[] { 1163 return this.mutableMessages 1164 } 1165 1166 getReadFileState(): FileStateCache { 1167 return this.readFileState 1168 } 1169 1170 getSessionId(): string { 1171 return getSessionId() 1172 } 1173 1174 setModel(model: string): void { 1175 this.config.userSpecifiedModel = model 1176 } 1177 } 1178 1179 /** 1180 * Sends a single prompt to the Claude API and returns the response. 1181 * Assumes that claude is being used non-interactively -- will not 1182 * ask the user for permissions or further input. 1183 * 1184 * Convenience wrapper around QueryEngine for one-shot usage. 1185 */ 1186 export async function* ask({ 1187 commands, 1188 prompt, 1189 promptUuid, 1190 isMeta, 1191 cwd, 1192 tools, 1193 mcpClients, 1194 verbose = false, 1195 thinkingConfig, 1196 maxTurns, 1197 maxBudgetUsd, 1198 taskBudget, 1199 canUseTool, 1200 mutableMessages = [], 1201 getReadFileCache, 1202 setReadFileCache, 1203 customSystemPrompt, 1204 appendSystemPrompt, 1205 userSpecifiedModel, 1206 fallbackModel, 1207 jsonSchema, 1208 getAppState, 1209 setAppState, 1210 abortController, 1211 replayUserMessages = false, 1212 includePartialMessages = false, 1213 handleElicitation, 1214 agents = [], 1215 setSDKStatus, 1216 orphanedPermission, 1217 }: { 1218 commands: Command[] 1219 prompt: string | Array<ContentBlockParam> 1220 promptUuid?: string 1221 isMeta?: boolean 1222 cwd: string 1223 tools: Tools 1224 verbose?: boolean 1225 mcpClients: MCPServerConnection[] 1226 thinkingConfig?: ThinkingConfig 1227 maxTurns?: number 1228 maxBudgetUsd?: number 1229 taskBudget?: { total: number } 1230 canUseTool: CanUseToolFn 1231 mutableMessages?: Message[] 1232 customSystemPrompt?: string 1233 appendSystemPrompt?: string 1234 userSpecifiedModel?: string 1235 fallbackModel?: string 1236 jsonSchema?: Record<string, unknown> 1237 getAppState: () => AppState 1238 setAppState: (f: (prev: AppState) => AppState) => void 1239 getReadFileCache: () => FileStateCache 1240 setReadFileCache: (cache: FileStateCache) => void 1241 abortController?: AbortController 1242 replayUserMessages?: boolean 1243 includePartialMessages?: boolean 1244 handleElicitation?: ToolUseContext['handleElicitation'] 1245 agents?: AgentDefinition[] 1246 setSDKStatus?: (status: SDKStatus) => void 1247 orphanedPermission?: OrphanedPermission 1248 }): AsyncGenerator<SDKMessage, void, unknown> { 1249 const engine = new QueryEngine({ 1250 cwd, 1251 tools, 1252 commands, 1253 mcpClients, 1254 agents, 1255 canUseTool, 1256 getAppState, 1257 setAppState, 1258 initialMessages: mutableMessages, 1259 readFileCache: cloneFileStateCache(getReadFileCache()), 1260 customSystemPrompt, 1261 appendSystemPrompt, 1262 userSpecifiedModel, 1263 fallbackModel, 1264 thinkingConfig, 1265 maxTurns, 1266 maxBudgetUsd, 1267 taskBudget, 1268 jsonSchema, 1269 verbose, 1270 handleElicitation, 1271 replayUserMessages, 1272 includePartialMessages, 1273 setSDKStatus, 1274 abortController, 1275 orphanedPermission, 1276 ...(feature('HISTORY_SNIP') 1277 ? { 1278 snipReplay: (yielded: Message, store: Message[]) => { 1279 if (!snipProjection!.isSnipBoundaryMessage(yielded)) 1280 return undefined 1281 return snipModule!.snipCompactIfNeeded(store, { force: true }) 1282 }, 1283 } 1284 : {}), 1285 }) 1286 1287 try { 1288 yield* engine.submitMessage(prompt, { 1289 uuid: promptUuid, 1290 isMeta, 1291 }) 1292 } finally { 1293 setReadFileCache(engine.getReadFileState()) 1294 } 1295 }