/ utils / queryContext.ts
queryContext.ts
  1  /**
  2   * Shared helpers for building the API cache-key prefix (systemPrompt,
  3   * userContext, systemContext) for query() calls.
  4   *
  5   * Lives in its own file because it imports from context.ts and
  6   * constants/prompts.ts, which are high in the dependency graph. Putting
  7   * these imports in systemPrompt.ts or sideQuestion.ts (both reachable
  8   * from commands.ts) would create cycles. Only entrypoint-layer files
  9   * import from here (QueryEngine.ts, cli/print.ts).
 10   */
 11  
 12  import type { Command } from '../commands.js'
 13  import { getSystemPrompt } from '../constants/prompts.js'
 14  import { getSystemContext, getUserContext } from '../context.js'
 15  import type { MCPServerConnection } from '../services/mcp/types.js'
 16  import type { AppState } from '../state/AppStateStore.js'
 17  import type { Tools, ToolUseContext } from '../Tool.js'
 18  import type { AgentDefinition } from '../tools/AgentTool/loadAgentsDir.js'
 19  import type { Message } from '../types/message.js'
 20  import { createAbortController } from './abortController.js'
 21  import type { FileStateCache } from './fileStateCache.js'
 22  import type { CacheSafeParams } from './forkedAgent.js'
 23  import { getMainLoopModel } from './model/model.js'
 24  import { asSystemPrompt } from './systemPromptType.js'
 25  import {
 26    shouldEnableThinkingByDefault,
 27    type ThinkingConfig,
 28  } from './thinking.js'
 29  
 30  /**
 31   * Fetch the three context pieces that form the API cache-key prefix:
 32   * systemPrompt parts, userContext, systemContext.
 33   *
 34   * When customSystemPrompt is set, the default getSystemPrompt build and
 35   * getSystemContext are skipped — the custom prompt replaces the default
 36   * entirely, and systemContext would be appended to a default that isn't
 37   * being used.
 38   *
 39   * Callers assemble the final systemPrompt from defaultSystemPrompt (or
 40   * customSystemPrompt) + optional extras + appendSystemPrompt. QueryEngine
 41   * injects coordinator userContext and memory-mechanics prompt on top;
 42   * sideQuestion's fallback uses the base result directly.
 43   */
 44  export async function fetchSystemPromptParts({
 45    tools,
 46    mainLoopModel,
 47    additionalWorkingDirectories,
 48    mcpClients,
 49    customSystemPrompt,
 50  }: {
 51    tools: Tools
 52    mainLoopModel: string
 53    additionalWorkingDirectories: string[]
 54    mcpClients: MCPServerConnection[]
 55    customSystemPrompt: string | undefined
 56  }): Promise<{
 57    defaultSystemPrompt: string[]
 58    userContext: { [k: string]: string }
 59    systemContext: { [k: string]: string }
 60  }> {
 61    const [defaultSystemPrompt, userContext, systemContext] = await Promise.all([
 62      customSystemPrompt !== undefined
 63        ? Promise.resolve([])
 64        : getSystemPrompt(
 65            tools,
 66            mainLoopModel,
 67            additionalWorkingDirectories,
 68            mcpClients,
 69          ),
 70      getUserContext(),
 71      customSystemPrompt !== undefined ? Promise.resolve({}) : getSystemContext(),
 72    ])
 73    return { defaultSystemPrompt, userContext, systemContext }
 74  }
 75  
 76  /**
 77   * Build CacheSafeParams from raw inputs when getLastCacheSafeParams() is null.
 78   *
 79   * Used by the SDK side_question handler (print.ts) on resume before a turn
 80   * completes — there's no stopHooks snapshot yet. Mirrors the system prompt
 81   * assembly in QueryEngine.ts:ask() so the rebuilt prefix matches what the
 82   * main loop will send, preserving the cache hit in the common case.
 83   *
 84   * May still miss the cache if the main loop applies extras this path doesn't
 85   * know about (coordinator mode, memory-mechanics prompt). That's acceptable —
 86   * the alternative is returning null and failing the side question entirely.
 87   */
 88  export async function buildSideQuestionFallbackParams({
 89    tools,
 90    commands,
 91    mcpClients,
 92    messages,
 93    readFileState,
 94    getAppState,
 95    setAppState,
 96    customSystemPrompt,
 97    appendSystemPrompt,
 98    thinkingConfig,
 99    agents,
100  }: {
101    tools: Tools
102    commands: Command[]
103    mcpClients: MCPServerConnection[]
104    messages: Message[]
105    readFileState: FileStateCache
106    getAppState: () => AppState
107    setAppState: (f: (prev: AppState) => AppState) => void
108    customSystemPrompt: string | undefined
109    appendSystemPrompt: string | undefined
110    thinkingConfig: ThinkingConfig | undefined
111    agents: AgentDefinition[]
112  }): Promise<CacheSafeParams> {
113    const mainLoopModel = getMainLoopModel()
114    const appState = getAppState()
115  
116    const { defaultSystemPrompt, userContext, systemContext } =
117      await fetchSystemPromptParts({
118        tools,
119        mainLoopModel,
120        additionalWorkingDirectories: Array.from(
121          appState.toolPermissionContext.additionalWorkingDirectories.keys(),
122        ),
123        mcpClients,
124        customSystemPrompt,
125      })
126  
127    const systemPrompt = asSystemPrompt([
128      ...(customSystemPrompt !== undefined
129        ? [customSystemPrompt]
130        : defaultSystemPrompt),
131      ...(appendSystemPrompt ? [appendSystemPrompt] : []),
132    ])
133  
134    // Strip in-progress assistant message (stop_reason === null) — same guard
135    // as btw.tsx. The SDK can fire side_question mid-turn.
136    const last = messages.at(-1)
137    const forkContextMessages =
138      last?.type === 'assistant' && last.message.stop_reason === null
139        ? messages.slice(0, -1)
140        : messages
141  
142    const toolUseContext: ToolUseContext = {
143      options: {
144        commands,
145        debug: false,
146        mainLoopModel,
147        tools,
148        verbose: false,
149        thinkingConfig:
150          thinkingConfig ??
151          (shouldEnableThinkingByDefault() !== false
152            ? { type: 'adaptive' }
153            : { type: 'disabled' }),
154        mcpClients,
155        mcpResources: {},
156        isNonInteractiveSession: true,
157        agentDefinitions: { activeAgents: agents, allAgents: [] },
158        customSystemPrompt,
159        appendSystemPrompt,
160      },
161      abortController: createAbortController(),
162      readFileState,
163      getAppState,
164      setAppState,
165      messages: forkContextMessages,
166      setInProgressToolUseIDs: () => {},
167      setResponseLength: () => {},
168      updateFileHistoryState: () => {},
169      updateAttributionState: () => {},
170    }
171  
172    return {
173      systemPrompt,
174      userContext,
175      systemContext,
176      toolUseContext,
177      forkContextMessages,
178    }
179  }