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 { AppState } from './state/AppState.js'
import { type Tools, type ToolUseContext, toolMatchesName } from './Tool.js' import { type Tools, type ToolUseContext, toolMatchesName } from './Tool.js'
import type { AgentDefinition } from '@claude-code-best/builtin-tools/tools/AgentTool/loadAgentsDir.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 { SYNTHETIC_OUTPUT_TOOL_NAME } from '@claude-code-best/builtin-tools/tools/SyntheticOutputTool/SyntheticOutputTool.js'
import type { APIError } from '@anthropic-ai/sdk' import type { APIError } from '@anthropic-ai/sdk'
import type { CompactMetadata, Message, SystemCompactBoundaryMessage } from './types/message.js' import type { CompactMetadata, Message, SystemCompactBoundaryMessage } from './types/message.js'
@ -77,7 +78,7 @@ import {
flushSessionStorage, flushSessionStorage,
recordTranscript, recordTranscript,
} from './utils/sessionStorage.js' } from './utils/sessionStorage.js'
import { asSystemPrompt } from './utils/systemPromptType.js' import { buildEffectiveSystemPrompt } from './utils/systemPrompt.js'
import { resolveThemeSetting } from './utils/systemTheme.js' import { resolveThemeSetting } from './utils/systemTheme.js'
import { import {
shouldEnableThinkingByDefault, shouldEnableThinkingByDefault,
@ -146,6 +147,7 @@ export type QueryEngineConfig = {
readFileCache: FileStateCache readFileCache: FileStateCache
customSystemPrompt?: string customSystemPrompt?: string
appendSystemPrompt?: string appendSystemPrompt?: string
mainThreadAgentDefinition?: AgentDefinition
userSpecifiedModel?: string userSpecifiedModel?: string
fallbackModel?: string fallbackModel?: string
thinkingConfig?: ThinkingConfig thinkingConfig?: ThinkingConfig
@ -229,6 +231,7 @@ export class QueryEngine {
canUseTool, canUseTool,
customSystemPrompt, customSystemPrompt,
appendSystemPrompt, appendSystemPrompt,
mainThreadAgentDefinition,
userSpecifiedModel, userSpecifiedModel,
fallbackModel, fallbackModel,
jsonSchema, jsonSchema,
@ -292,6 +295,13 @@ export class QueryEngine {
// Narrow once so TS tracks the type through the conditionals below. // Narrow once so TS tracks the type through the conditionals below.
const customPrompt = const customPrompt =
typeof customSystemPrompt === 'string' ? customSystemPrompt : undefined 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 { const {
defaultSystemPrompt, defaultSystemPrompt,
userContext: baseUserContext, userContext: baseUserContext,
@ -303,7 +313,7 @@ export class QueryEngine {
initialAppState.toolPermissionContext.additionalWorkingDirectories.keys(), initialAppState.toolPermissionContext.additionalWorkingDirectories.keys(),
), ),
mcpClients, mcpClients,
customSystemPrompt: customPrompt, customSystemPrompt: customPromptForFetch,
}) })
headlessProfilerCheckpoint('after_getSystemPrompt') headlessProfilerCheckpoint('after_getSystemPrompt')
const userContext = { const userContext = {
@ -325,11 +335,18 @@ export class QueryEngine {
? await loadMemoryPrompt() ? await loadMemoryPrompt()
: null : null
const systemPrompt = asSystemPrompt([ const combinedAppendSystemPrompt = [
...(customPrompt !== undefined ? [customPrompt] : defaultSystemPrompt),
...(memoryMechanicsPrompt ? [memoryMechanicsPrompt] : []), ...(memoryMechanicsPrompt ? [memoryMechanicsPrompt] : []),
...(appendSystemPrompt ? [appendSystemPrompt] : []), ...(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 // Register function hook for structured output enforcement
const hasStructuredOutputTool = tools.some(t => const hasStructuredOutputTool = tools.some(t =>
@ -1245,6 +1262,7 @@ export async function* ask({
setReadFileCache, setReadFileCache,
customSystemPrompt, customSystemPrompt,
appendSystemPrompt, appendSystemPrompt,
mainThreadAgentDefinition,
userSpecifiedModel, userSpecifiedModel,
fallbackModel, fallbackModel,
jsonSchema, jsonSchema,
@ -1274,6 +1292,7 @@ export async function* ask({
mutableMessages?: Message[] mutableMessages?: Message[]
customSystemPrompt?: string customSystemPrompt?: string
appendSystemPrompt?: string appendSystemPrompt?: string
mainThreadAgentDefinition?: AgentDefinition
userSpecifiedModel?: string userSpecifiedModel?: string
fallbackModel?: string fallbackModel?: string
jsonSchema?: Record<string, unknown> jsonSchema?: Record<string, unknown>
@ -1302,6 +1321,7 @@ export async function* ask({
readFileCache: cloneFileStateCache(getReadFileCache()), readFileCache: cloneFileStateCache(getReadFileCache()),
customSystemPrompt, customSystemPrompt,
appendSystemPrompt, appendSystemPrompt,
mainThreadAgentDefinition,
userSpecifiedModel, userSpecifiedModel,
fallbackModel, fallbackModel,
thinkingConfig, thinkingConfig,

View File

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