feat(query-engine): pass mainThreadAgentDefinition to support built-in agent system prompts

This commit is contained in:
yhangf 2026-05-08 14:43:11 +08:00
parent 3f84d157e5
commit 697ee49752
2 changed files with 28 additions and 5 deletions

View File

@ -39,6 +39,7 @@ import type { MCPServerConnection } from './services/mcp/types.js'
import type { AppState } from './state/AppState.js'
import { type Tools, type ToolUseContext, toolMatchesName } from './Tool.js'
import type { AgentDefinition } from '@claude-code-best/builtin-tools/tools/AgentTool/loadAgentsDir.js'
import { isBuiltInAgent } from '@claude-code-best/builtin-tools/tools/AgentTool/loadAgentsDir.js'
import { SYNTHETIC_OUTPUT_TOOL_NAME } from '@claude-code-best/builtin-tools/tools/SyntheticOutputTool/SyntheticOutputTool.js'
import type { APIError } from '@anthropic-ai/sdk'
import type { CompactMetadata, Message, SystemCompactBoundaryMessage } from './types/message.js'
@ -77,7 +78,7 @@ import {
flushSessionStorage,
recordTranscript,
} from './utils/sessionStorage.js'
import { asSystemPrompt } from './utils/systemPromptType.js'
import { buildEffectiveSystemPrompt } from './utils/systemPrompt.js'
import { resolveThemeSetting } from './utils/systemTheme.js'
import {
shouldEnableThinkingByDefault,
@ -146,6 +147,7 @@ export type QueryEngineConfig = {
readFileCache: FileStateCache
customSystemPrompt?: string
appendSystemPrompt?: string
mainThreadAgentDefinition?: AgentDefinition
userSpecifiedModel?: string
fallbackModel?: string
thinkingConfig?: ThinkingConfig
@ -229,6 +231,7 @@ export class QueryEngine {
canUseTool,
customSystemPrompt,
appendSystemPrompt,
mainThreadAgentDefinition,
userSpecifiedModel,
fallbackModel,
jsonSchema,
@ -292,6 +295,13 @@ export class QueryEngine {
// Narrow once so TS tracks the type through the conditionals below.
const customPrompt =
typeof customSystemPrompt === 'string' ? customSystemPrompt : undefined
// For built-in agents, don't pass customSystemPrompt to fetchSystemPromptParts
// (it's only used for cache key / context building there). The actual system
// prompt is assembled below via buildEffectiveSystemPrompt.
const customPromptForFetch =
mainThreadAgentDefinition && isBuiltInAgent(mainThreadAgentDefinition)
? undefined
: customPrompt
const {
defaultSystemPrompt,
userContext: baseUserContext,
@ -303,7 +313,7 @@ export class QueryEngine {
initialAppState.toolPermissionContext.additionalWorkingDirectories.keys(),
),
mcpClients,
customSystemPrompt: customPrompt,
customSystemPrompt: customPromptForFetch,
})
headlessProfilerCheckpoint('after_getSystemPrompt')
const userContext = {
@ -325,11 +335,18 @@ export class QueryEngine {
? await loadMemoryPrompt()
: null
const systemPrompt = asSystemPrompt([
...(customPrompt !== undefined ? [customPrompt] : defaultSystemPrompt),
const combinedAppendSystemPrompt = [
...(memoryMechanicsPrompt ? [memoryMechanicsPrompt] : []),
...(appendSystemPrompt ? [appendSystemPrompt] : []),
])
].join('\n') || undefined
const systemPrompt = buildEffectiveSystemPrompt({
mainThreadAgentDefinition,
toolUseContext: { options: { customSystemPrompt: customPrompt } as ToolUseContext['options'] },
customSystemPrompt: customPrompt,
defaultSystemPrompt,
appendSystemPrompt: combinedAppendSystemPrompt,
})
// Register function hook for structured output enforcement
const hasStructuredOutputTool = tools.some(t =>
@ -1245,6 +1262,7 @@ export async function* ask({
setReadFileCache,
customSystemPrompt,
appendSystemPrompt,
mainThreadAgentDefinition,
userSpecifiedModel,
fallbackModel,
jsonSchema,
@ -1274,6 +1292,7 @@ export async function* ask({
mutableMessages?: Message[]
customSystemPrompt?: string
appendSystemPrompt?: string
mainThreadAgentDefinition?: AgentDefinition
userSpecifiedModel?: string
fallbackModel?: string
jsonSchema?: Record<string, unknown>
@ -1302,6 +1321,7 @@ export async function* ask({
readFileCache: cloneFileStateCache(getReadFileCache()),
customSystemPrompt,
appendSystemPrompt,
mainThreadAgentDefinition,
userSpecifiedModel,
fallbackModel,
thinkingConfig,

View File

@ -2217,6 +2217,9 @@ function runHeadlessStreaming(
},
customSystemPrompt: options.systemPrompt,
appendSystemPrompt: options.appendSystemPrompt,
mainThreadAgentDefinition: currentAgents.find(
a => a.agentType === getMainThreadAgentType(),
),
getAppState,
setAppState,
abortController,