feat(agents): propagate allowedAgentTypes from agent tool resolution to query engine

This commit is contained in:
yhangf 2026-05-09 15:15:10 +08:00
parent 1efba6b17a
commit 396e4cdbe3
2 changed files with 13 additions and 5 deletions

View File

@ -140,6 +140,7 @@ export type QueryEngineConfig = {
commands: Command[] commands: Command[]
mcpClients: MCPServerConnection[] mcpClients: MCPServerConnection[]
agents: AgentDefinition[] agents: AgentDefinition[]
allowedAgentTypes?: string[]
canUseTool: CanUseToolFn canUseTool: CanUseToolFn
getAppState: () => AppState getAppState: () => AppState
setAppState: (f: (prev: AppState) => AppState) => void setAppState: (f: (prev: AppState) => AppState) => void
@ -240,6 +241,7 @@ export class QueryEngine {
replayUserMessages = false, replayUserMessages = false,
includePartialMessages = false, includePartialMessages = false,
agents = [], agents = [],
allowedAgentTypes,
setSDKStatus, setSDKStatus,
orphanedPermission, orphanedPermission,
} = this.config } = this.config
@ -383,7 +385,7 @@ export class QueryEngine {
isNonInteractiveSession: true, isNonInteractiveSession: true,
customSystemPrompt, customSystemPrompt,
appendSystemPrompt, appendSystemPrompt,
agentDefinitions: { activeAgents: agents, allAgents: [] }, agentDefinitions: { activeAgents: agents, allAgents: [], ...(allowedAgentTypes ? { allowedAgentTypes } : {}) },
theme: resolveThemeSetting(getGlobalConfig().theme), theme: resolveThemeSetting(getGlobalConfig().theme),
maxBudgetUsd, maxBudgetUsd,
}, },
@ -533,7 +535,7 @@ export class QueryEngine {
customSystemPrompt, customSystemPrompt,
appendSystemPrompt, appendSystemPrompt,
theme: resolveThemeSetting(getGlobalConfig().theme), theme: resolveThemeSetting(getGlobalConfig().theme),
agentDefinitions: { activeAgents: agents, allAgents: [] }, agentDefinitions: { activeAgents: agents, allAgents: [], ...(allowedAgentTypes ? { allowedAgentTypes } : {}) },
maxBudgetUsd, maxBudgetUsd,
}, },
getAppState, getAppState,
@ -1273,6 +1275,7 @@ export async function* ask({
includePartialMessages = false, includePartialMessages = false,
handleElicitation, handleElicitation,
agents = [], agents = [],
allowedAgentTypes,
setSDKStatus, setSDKStatus,
orphanedPermission, orphanedPermission,
}: { }: {
@ -1305,6 +1308,7 @@ export async function* ask({
includePartialMessages?: boolean includePartialMessages?: boolean
handleElicitation?: ToolUseContext['handleElicitation'] handleElicitation?: ToolUseContext['handleElicitation']
agents?: AgentDefinition[] agents?: AgentDefinition[]
allowedAgentTypes?: string[]
setSDKStatus?: (status: SDKStatus) => void setSDKStatus?: (status: SDKStatus) => void
orphanedPermission?: OrphanedPermission orphanedPermission?: OrphanedPermission
}): AsyncGenerator<SDKMessage, void, unknown> { }): AsyncGenerator<SDKMessage, void, unknown> {
@ -1314,6 +1318,7 @@ export async function* ask({
commands, commands,
mcpClients, mcpClients,
agents: agents ?? [], agents: agents ?? [],
allowedAgentTypes,
canUseTool, canUseTool,
getAppState, getAppState,
setAppState, setAppState,

View File

@ -2032,9 +2032,11 @@ function runHeadlessStreaming(
const mainThreadAgentDef = currentAgents.find( const mainThreadAgentDef = currentAgents.find(
a => a.agentType === getMainThreadAgentType(), a => a.agentType === getMainThreadAgentType(),
) )
const allTools = mainThreadAgentDef const resolvedAgentResult = mainThreadAgentDef
? resolveAgentTools(mainThreadAgentDef, rawTools, false, true).resolvedTools ? resolveAgentTools(mainThreadAgentDef, rawTools, false, true)
: rawTools : null
const allTools = resolvedAgentResult ? resolvedAgentResult.resolvedTools : rawTools
const mainThreadAllowedAgentTypes = resolvedAgentResult?.allowedAgentTypes
for (const uuid of batchUuids) { for (const uuid of batchUuids) {
notifyCommandLifecycle(uuid, 'started') notifyCommandLifecycle(uuid, 'started')
@ -2245,6 +2247,7 @@ function runHeadlessStreaming(
: undefined, : undefined,
), ),
agents: currentAgents, agents: currentAgents,
allowedAgentTypes: mainThreadAllowedAgentTypes,
orphanedPermission: cmd.orphanedPermission, orphanedPermission: cmd.orphanedPermission,
setSDKStatus: status => { setSDKStatus: status => {
output.enqueue({ output.enqueue({