From 697ee49752d64b03b19b8864d53da2a81e19efcf Mon Sep 17 00:00:00 2001
From: yhangf <yhf5fhy@gmail.com>
Date: Fri, 8 May 2026 14:43:11 +0800
Subject: [PATCH] feat(query-engine): pass mainThreadAgentDefinition to support
 built-in agent system prompts

---
 src/QueryEngine.ts | 30 +++++++++++++++++++++++++-----
 src/cli/print.ts   |  3 +++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/QueryEngine.ts b/src/QueryEngine.ts
index db5732489..69b2c1415 100644
--- a/src/QueryEngine.ts
+++ b/src/QueryEngine.ts
@@ -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,
diff --git a/src/cli/print.ts b/src/cli/print.ts
index 5d3a1ed3c..71ff0d68b 100644
--- a/src/cli/print.ts
+++ b/src/cli/print.ts
@@ -2217,6 +2217,9 @@ function runHeadlessStreaming(
                   },
                   customSystemPrompt: options.systemPrompt,
                   appendSystemPrompt: options.appendSystemPrompt,
+                  mainThreadAgentDefinition: currentAgents.find(
+                    a => a.agentType === getMainThreadAgentType(),
+                  ),
                   getAppState,
                   setAppState,
                   abortController,