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
This commit is contained in:
Taylor Mullen
2026-04-01 15:42:23 -07:00
parent cce060646c
commit 5aba28c8be
6 changed files with 130 additions and 15 deletions
+2 -1
View File
@@ -33,6 +33,7 @@ export class SubagentTool extends BaseDeclarativeTool<AgentInputs, ToolResult> {
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<AgentInputs, ToolResult> {
super(
definition.name,
definition.displayName ?? definition.name,
definition.description,
descriptionOverride ?? definition.description,
Kind.Agent,
inputSchema,
messageBus,
+19 -13
View File
@@ -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<void> {
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<Promise<void>> = [];
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);
}
/**
+16 -1
View File
@@ -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
@@ -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.",
);
});
});
@@ -140,6 +140,7 @@ export class PromptProvider {
contextFilenames,
topicUpdateNarration: context.config.isTopicUpdateNarrationEnabled(),
})),
activeTeam: context.config.getActiveTeam(),
subAgents: this.withSection('agentContexts', () =>
context.config
.getAgentRegistry()
+14
View File
@@ -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