From 396e4cdbe34016fd7c58bd34d4de4a4ee8c5ed4a Mon Sep 17 00:00:00 2001
From: yhangf <yhf5fhy@gmail.com>
Date: Sat, 9 May 2026 15:15:10 +0800
Subject: [PATCH] feat(agents): propagate allowedAgentTypes from agent tool
 resolution to query engine

---
 src/QueryEngine.ts | 9 +++++++--
 src/cli/print.ts   | 9 ++++++---
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/QueryEngine.ts b/src/QueryEngine.ts
index 69b2c1415..3a78d19ba 100644
--- a/src/QueryEngine.ts
+++ b/src/QueryEngine.ts
@@ -140,6 +140,7 @@ export type QueryEngineConfig = {
   commands: Command[]
   mcpClients: MCPServerConnection[]
   agents: AgentDefinition[]
+  allowedAgentTypes?: string[]
   canUseTool: CanUseToolFn
   getAppState: () => AppState
   setAppState: (f: (prev: AppState) => AppState) => void
@@ -240,6 +241,7 @@ export class QueryEngine {
       replayUserMessages = false,
       includePartialMessages = false,
       agents = [],
+      allowedAgentTypes,
       setSDKStatus,
       orphanedPermission,
     } = this.config
@@ -383,7 +385,7 @@ export class QueryEngine {
         isNonInteractiveSession: true,
         customSystemPrompt,
         appendSystemPrompt,
-        agentDefinitions: { activeAgents: agents, allAgents: [] },
+        agentDefinitions: { activeAgents: agents, allAgents: [], ...(allowedAgentTypes ? { allowedAgentTypes } : {}) },
         theme: resolveThemeSetting(getGlobalConfig().theme),
         maxBudgetUsd,
       },
@@ -533,7 +535,7 @@ export class QueryEngine {
         customSystemPrompt,
         appendSystemPrompt,
         theme: resolveThemeSetting(getGlobalConfig().theme),
-        agentDefinitions: { activeAgents: agents, allAgents: [] },
+        agentDefinitions: { activeAgents: agents, allAgents: [], ...(allowedAgentTypes ? { allowedAgentTypes } : {}) },
         maxBudgetUsd,
       },
       getAppState,
@@ -1273,6 +1275,7 @@ export async function* ask({
   includePartialMessages = false,
   handleElicitation,
   agents = [],
+  allowedAgentTypes,
   setSDKStatus,
   orphanedPermission,
 }: {
@@ -1305,6 +1308,7 @@ export async function* ask({
   includePartialMessages?: boolean
   handleElicitation?: ToolUseContext['handleElicitation']
   agents?: AgentDefinition[]
+  allowedAgentTypes?: string[]
   setSDKStatus?: (status: SDKStatus) => void
   orphanedPermission?: OrphanedPermission
 }): AsyncGenerator<SDKMessage, void, unknown> {
@@ -1314,6 +1318,7 @@ export async function* ask({
     commands,
     mcpClients,
     agents: agents ?? [],
+    allowedAgentTypes,
     canUseTool,
     getAppState,
     setAppState,
diff --git a/src/cli/print.ts b/src/cli/print.ts
index f3e6f996c..67264c1aa 100644
--- a/src/cli/print.ts
+++ b/src/cli/print.ts
@@ -2032,9 +2032,11 @@ function runHeadlessStreaming(
           const mainThreadAgentDef = currentAgents.find(
             a => a.agentType === getMainThreadAgentType(),
           )
-          const allTools = mainThreadAgentDef
-            ? resolveAgentTools(mainThreadAgentDef, rawTools, false, true).resolvedTools
-            : rawTools
+          const resolvedAgentResult = mainThreadAgentDef
+            ? resolveAgentTools(mainThreadAgentDef, rawTools, false, true)
+            : null
+          const allTools = resolvedAgentResult ? resolvedAgentResult.resolvedTools : rawTools
+          const mainThreadAllowedAgentTypes = resolvedAgentResult?.allowedAgentTypes
 
           for (const uuid of batchUuids) {
             notifyCommandLifecycle(uuid, 'started')
@@ -2245,6 +2247,7 @@ function runHeadlessStreaming(
                         : undefined,
                     ),
                   agents: currentAgents,
+                  allowedAgentTypes: mainThreadAllowedAgentTypes,
                   orphanedPermission: cmd.orphanedPermission,
                   setSDKStatus: status => {
                     output.enqueue({