From 5aba28c8beedde8ebeebb7d5fa72e372412c63d8 Mon Sep 17 00:00:00 2001 From: Taylor Mullen Date: Wed, 1 Apr 2026 15:42:23 -0700 Subject: [PATCH] feat(core): enable team-aware orchestration in top-level agent - Update SystemPromptOptions to include activeTeam - Implement renderActiveTeam snippet for team instructions injection - Update PromptProvider to pass activeTeam from config to system prompt - Enhance SubagentTool to support description overrides - Update Config to prioritize and label team-based subagent tools - Ensure synchronized team reload in TeamRegistry and Config - Add unit tests for team-aware prompt generation --- packages/core/src/agents/subagent-tool.ts | 3 +- packages/core/src/agents/teamRegistry.ts | 32 ++++---- packages/core/src/config/config.ts | 17 +++- .../src/prompts/promptProvider-teams.test.ts | 78 +++++++++++++++++++ packages/core/src/prompts/promptProvider.ts | 1 + packages/core/src/prompts/snippets.ts | 14 ++++ 6 files changed, 130 insertions(+), 15 deletions(-) create mode 100644 packages/core/src/prompts/promptProvider-teams.test.ts diff --git a/packages/core/src/agents/subagent-tool.ts b/packages/core/src/agents/subagent-tool.ts index 3ef9f0aa86..f29acb9f3b 100644 --- a/packages/core/src/agents/subagent-tool.ts +++ b/packages/core/src/agents/subagent-tool.ts @@ -33,6 +33,7 @@ export class SubagentTool extends BaseDeclarativeTool { private readonly definition: AgentDefinition, private readonly context: AgentLoopContext, messageBus: MessageBus, + descriptionOverride?: string, ) { const inputSchema = definition.inputConfig.inputSchema; @@ -47,7 +48,7 @@ export class SubagentTool extends BaseDeclarativeTool { super( definition.name, definition.displayName ?? definition.name, - definition.description, + descriptionOverride ?? definition.description, Kind.Agent, inputSchema, messageBus, diff --git a/packages/core/src/agents/teamRegistry.ts b/packages/core/src/agents/teamRegistry.ts index 1157a4bc8d..97185619b2 100644 --- a/packages/core/src/agents/teamRegistry.ts +++ b/packages/core/src/agents/teamRegistry.ts @@ -48,7 +48,7 @@ export class TeamRegistry { // Load user-level teams first const userTeamsDir = Storage.getUserTeamsDir(); const userLoadResult = await loadTeamsFromDirectory(userTeamsDir); - this.processLoadResult(userLoadResult); + await this.processLoadResult(userLoadResult); const folderTrustEnabled = this.config.getFolderTrust(); const isTrustedFolder = this.config.isTrustedFolder(); @@ -65,7 +65,7 @@ export class TeamRegistry { // Load project-level teams (takes precedence over user-level if names collide) const projectTeamsDir = this.config.storage.getProjectTeamsDir(); const projectLoadResult = await loadTeamsFromDirectory(projectTeamsDir); - this.processLoadResult(projectLoadResult); + await this.processLoadResult(projectLoadResult); } if (this.config.getDebugMode()) { @@ -73,29 +73,35 @@ export class TeamRegistry { } } - private processLoadResult(result: TeamLoadResult): void { + private async processLoadResult(result: TeamLoadResult): Promise { for (const error of result.errors) { debugLogger.warn(`[TeamRegistry] Error loading team: ${error.message}`); coreEvents.emitFeedback('error', `Team loading error: ${error.message}`); } + const registrationPromises: Array> = []; + for (const team of result.teams) { this.teams.set(team.name, team); // Register team agents in the global AgentRegistry so they are available as SubagentTools for (const agent of team.agents) { - this.agentRegistry.registerAgent(agent).catch((e) => { - debugLogger.warn( - `[TeamRegistry] Error registering agent "${agent.name}" from team "${team.name}":`, - e, - ); - coreEvents.emitFeedback( - 'error', - `Error registering agent "${agent.name}" from team "${team.name}": ${e instanceof Error ? e.message : String(e)}`, - ); - }); + registrationPromises.push( + this.agentRegistry.registerAgent(agent).catch((e) => { + debugLogger.warn( + `[TeamRegistry] Error registering agent "${agent.name}" from team "${team.name}":`, + e, + ); + coreEvents.emitFeedback( + 'error', + `Error registering agent "${agent.name}" from team "${team.name}": ${e instanceof Error ? e.message : String(e)}`, + ); + }), + ); } } + + await Promise.allSettled(registrationPromises); } /** diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 6c6695d10a..617f4c7dbd 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -3617,6 +3617,9 @@ export class Config implements McpContext, AgentLoopContext { } const discoveredNames = this.agentRegistry.getAllDiscoveredAgentNames(); + const activeTeam = this.teamRegistry.getActiveTeam(); + const teamAgentNames = new Set(activeTeam?.agents.map((a) => a.name) ?? []); + for (const agentName of discoveredNames) { const definition = this.agentRegistry.getDiscoveredDefinition(agentName); if (!definition) { @@ -3630,7 +3633,17 @@ export class Config implements McpContext, AgentLoopContext { continue; } - const tool = new SubagentTool(definition, this, this.messageBus); + let descriptionOverride = definition.description; + if (teamAgentNames.has(definition.name)) { + descriptionOverride = `(Team Agent: ${activeTeam?.displayName}) ${definition.description}`; + } + + const tool = new SubagentTool( + definition, + this, + this.messageBus, + descriptionOverride, + ); registry.registerTool(tool); } catch (e: unknown) { debugLogger.warn( @@ -3730,8 +3743,10 @@ export class Config implements McpContext, AgentLoopContext { } private onAgentsRefreshed = async () => { + debugLogger.log('[Config] onAgentsRefreshed triggered'); await this.teamRegistry.reload(); if (this._toolRegistry) { + debugLogger.log('[Config] Re-registering sub-agent tools'); this.registerSubAgentTools(this._toolRegistry); } // Propagate updates to the active chat session diff --git a/packages/core/src/prompts/promptProvider-teams.test.ts b/packages/core/src/prompts/promptProvider-teams.test.ts new file mode 100644 index 0000000000..1e3c47e30a --- /dev/null +++ b/packages/core/src/prompts/promptProvider-teams.test.ts @@ -0,0 +1,78 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { PromptProvider } from './promptProvider.js'; +import { type AgentLoopContext } from '../config/agent-loop-context.js'; +import { type TeamDefinition } from '../agents/types.js'; + +describe('PromptProvider with Agent Teams', () => { + let promptProvider: PromptProvider; + let mockContext: AgentLoopContext; + + beforeEach(() => { + promptProvider = new PromptProvider(); + mockContext = { + config: { + isInteractive: vi.fn().mockReturnValue(true), + getSkillManager: vi.fn().mockReturnValue({ getSkills: () => [] }), + getAgentRegistry: vi + .fn() + .mockReturnValue({ getAllDefinitions: () => [] }), + getActiveModel: vi.fn().mockReturnValue('gemini-3.1-pro-preview'), + getActiveTeam: vi.fn().mockReturnValue(undefined), + getApprovedPlanPath: vi.fn().mockReturnValue(undefined), + getApprovalMode: vi.fn().mockReturnValue('default'), + getGemini31LaunchedSync: vi.fn().mockReturnValue(true), + getGemini31FlashLiteLaunchedSync: vi.fn().mockReturnValue(true), + getHasAccessToPreviewModel: vi.fn().mockReturnValue(true), + isTrackerEnabled: vi.fn().mockReturnValue(false), + isTopicUpdateNarrationEnabled: vi.fn().mockReturnValue(false), + isMemoryManagerEnabled: vi.fn().mockReturnValue(false), + isInteractiveShellEnabled: vi.fn().mockReturnValue(true), + getEnableShellOutputEfficiency: vi.fn().mockReturnValue(true), + getSandboxEnabled: vi.fn().mockReturnValue(false), + storage: { + getPlansDir: vi.fn().mockReturnValue('/tmp/plans'), + }, + topicState: { + getTopic: vi.fn().mockReturnValue(undefined), + }, + }, + toolRegistry: { + getAllToolNames: vi.fn().mockReturnValue([]), + getAllTools: vi.fn().mockReturnValue([]), + }, + } as unknown as AgentLoopContext; + }); + + it('should not include team section when no team is active', () => { + const prompt = promptProvider.getCoreSystemPrompt( + mockContext as unknown as AgentLoopContext, + ); + expect(prompt).not.toContain('# Active Agent Team'); + }); + + it('should include team section when a team is active', () => { + const mockTeam: TeamDefinition = { + name: 'test-team', + displayName: 'Test Team', + description: 'A test team', + instructions: 'These are the team instructions.', + agents: [], + }; + mockContext.config.getActiveTeam.mockReturnValue(mockTeam); + + const prompt = promptProvider.getCoreSystemPrompt( + mockContext as unknown as AgentLoopContext, + ); + expect(prompt).toContain('# Active Agent Team: Test Team'); + expect(prompt).toContain('These are the team instructions.'); + expect(prompt).toContain( + "You should prioritize delegating tasks to this team's agents whenever appropriate.", + ); + }); +}); diff --git a/packages/core/src/prompts/promptProvider.ts b/packages/core/src/prompts/promptProvider.ts index 0036dae560..e0b6105cbc 100644 --- a/packages/core/src/prompts/promptProvider.ts +++ b/packages/core/src/prompts/promptProvider.ts @@ -140,6 +140,7 @@ export class PromptProvider { contextFilenames, topicUpdateNarration: context.config.isTopicUpdateNarrationEnabled(), })), + activeTeam: context.config.getActiveTeam(), subAgents: this.withSection('agentContexts', () => context.config .getAgentRegistry() diff --git a/packages/core/src/prompts/snippets.ts b/packages/core/src/prompts/snippets.ts index 5440583419..5440d354a6 100644 --- a/packages/core/src/prompts/snippets.ts +++ b/packages/core/src/prompts/snippets.ts @@ -36,12 +36,14 @@ import { } from '../tools/tool-names.js'; import type { HierarchicalMemory } from '../config/memory.js'; import { DEFAULT_CONTEXT_FILENAME } from '../tools/memoryTool.js'; +import { type TeamDefinition } from '../agents/types.js'; // --- Options Structs --- export interface SystemPromptOptions { preamble?: PreambleOptions; coreMandates?: CoreMandatesOptions; + activeTeam?: TeamDefinition; subAgents?: SubAgentOptions[]; agentSkills?: AgentSkillOptions[]; hookContext?: boolean; @@ -126,6 +128,8 @@ ${renderPreamble(options.preamble)} ${renderCoreMandates(options.coreMandates)} +${renderActiveTeam(options.activeTeam)} + ${renderSubAgents(options.subAgents)} ${renderAgentSkills(options.agentSkills)} @@ -247,6 +251,16 @@ Use the following guidelines to optimize your search and read patterns. `.trim(); } +export function renderActiveTeam(team?: TeamDefinition): string { + if (!team) return ''; + return ` +# Active Agent Team: ${team.displayName} + +${team.instructions} + +You should prioritize delegating tasks to this team's agents whenever appropriate.`.trim(); +} + export function renderSubAgents(subAgents?: SubAgentOptions[]): string { if (!subAgents || subAgents.length === 0) return ''; const subAgentsXml = subAgents