feat(agents): Add remote agents to agent registry (#15711)

2026-03-18 01:51:20 -07:00 · 2025-12-30 18:11:51 -05:00
parent 0eb84f5133
commit 3ebe4e6a8f
4 changed files with 195 additions and 90 deletions
--- a/packages/core/src/agents/registry.test.ts
+++ b/packages/core/src/agents/registry.test.ts
@@ -11,6 +11,7 @@ import type { AgentDefinition, LocalAgentDefinition } from './types.js';
 import type { Config } from '../config/config.js';
 import { debugLogger } from '../utils/debugLogger.js';
 import { coreEvents, CoreEvent } from '../utils/events.js';
+import { A2AClientManager } from './a2a-client-manager.js';
 import {
  DEFAULT_GEMINI_FLASH_LITE_MODEL,
  GEMINI_MODEL_ALIAS_AUTO,
@@ -26,10 +27,16 @@ vi.mock('./toml-loader.js', () => ({
    .mockResolvedValue({ agents: [], errors: [] }),
 }));

+vi.mock('./a2a-client-manager.js', () => ({
+  A2AClientManager: {
+    getInstance: vi.fn(),
+  },
+}));
+
 // A test-only subclass to expose the protected `registerAgent` method.
 class TestableAgentRegistry extends AgentRegistry {
-  testRegisterAgent(definition: AgentDefinition): void {
-    this.registerAgent(definition);
+  async testRegisterAgent(definition: AgentDefinition): Promise<void> {
+    await this.registerAgent(definition);
  }
 }

@@ -237,8 +244,8 @@ describe('AgentRegistry', () => {
  });

  describe('registration logic', () => {
-    it('should register a valid agent definition', () => {
-      registry.testRegisterAgent(MOCK_AGENT_V1);
+    it('should register a valid agent definition', async () => {
+      await registry.testRegisterAgent(MOCK_AGENT_V1);
      expect(registry.getDefinition('MockAgent')).toEqual(MOCK_AGENT_V1);
      expect(
        mockConfig.modelConfigService.getResolvedConfig({
@@ -257,7 +264,7 @@ describe('AgentRegistry', () => {
      });
    });

-    it('should register a remote agent definition', () => {
+    it('should register a remote agent definition', async () => {
      const remoteAgent: AgentDefinition = {
        kind: 'remote',
        name: 'RemoteAgent',
@@ -265,11 +272,16 @@ describe('AgentRegistry', () => {
        agentCardUrl: 'https://example.com/card',
        inputConfig: { inputs: {} },
      };
-      registry.testRegisterAgent(remoteAgent);
+
+      vi.mocked(A2AClientManager.getInstance).mockReturnValue({
+        loadAgent: vi.fn().mockResolvedValue({ name: 'RemoteAgent' }),
+      } as unknown as A2AClientManager);
+
+      await registry.testRegisterAgent(remoteAgent);
      expect(registry.getDefinition('RemoteAgent')).toEqual(remoteAgent);
    });

-    it('should log remote agent registration in debug mode', () => {
+    it('should log remote agent registration in debug mode', async () => {
      const debugConfig = makeFakeConfig({ debugMode: true });
      const debugRegistry = new TestableAgentRegistry(debugConfig);
      const debugLogSpy = vi
@@ -284,31 +296,35 @@ describe('AgentRegistry', () => {
        inputConfig: { inputs: {} },
      };

-      debugRegistry.testRegisterAgent(remoteAgent);
+      vi.mocked(A2AClientManager.getInstance).mockReturnValue({
+        loadAgent: vi.fn().mockResolvedValue({ name: 'RemoteAgent' }),
+      } as unknown as A2AClientManager);
+
+      await debugRegistry.testRegisterAgent(remoteAgent);

      expect(debugLogSpy).toHaveBeenCalledWith(
        `[AgentRegistry] Registered remote agent 'RemoteAgent' with card: https://example.com/card`,
      );
    });

-    it('should handle special characters in agent names', () => {
+    it('should handle special characters in agent names', async () => {
      const specialAgent = {
        ...MOCK_AGENT_V1,
        name: 'Agent-123_$pecial.v2',
      };
-      registry.testRegisterAgent(specialAgent);
+      await registry.testRegisterAgent(specialAgent);
      expect(registry.getDefinition('Agent-123_$pecial.v2')).toEqual(
        specialAgent,
      );
    });

-    it('should reject an agent definition missing a name', () => {
+    it('should reject an agent definition missing a name', async () => {
      const invalidAgent = { ...MOCK_AGENT_V1, name: '' };
      const debugWarnSpy = vi
        .spyOn(debugLogger, 'warn')
        .mockImplementation(() => {});

-      registry.testRegisterAgent(invalidAgent);
+      await registry.testRegisterAgent(invalidAgent);

      expect(registry.getDefinition('MockAgent')).toBeUndefined();
      expect(debugWarnSpy).toHaveBeenCalledWith(
@@ -316,13 +332,13 @@ describe('AgentRegistry', () => {
      );
    });

-    it('should reject an agent definition missing a description', () => {
+    it('should reject an agent definition missing a description', async () => {
      const invalidAgent = { ...MOCK_AGENT_V1, description: '' };
      const debugWarnSpy = vi
        .spyOn(debugLogger, 'warn')
        .mockImplementation(() => {});

-      registry.testRegisterAgent(invalidAgent as AgentDefinition);
+      await registry.testRegisterAgent(invalidAgent as AgentDefinition);

      expect(registry.getDefinition('MockAgent')).toBeUndefined();
      expect(debugWarnSpy).toHaveBeenCalledWith(
@@ -330,41 +346,41 @@ describe('AgentRegistry', () => {
      );
    });

-    it('should overwrite an existing agent definition', () => {
-      registry.testRegisterAgent(MOCK_AGENT_V1);
+    it('should overwrite an existing agent definition', async () => {
+      await registry.testRegisterAgent(MOCK_AGENT_V1);
      expect(registry.getDefinition('MockAgent')?.description).toBe(
        'Mock Description V1',
      );

-      registry.testRegisterAgent(MOCK_AGENT_V2);
+      await registry.testRegisterAgent(MOCK_AGENT_V2);
      expect(registry.getDefinition('MockAgent')?.description).toBe(
        'Mock Description V2 (Updated)',
      );
      expect(registry.getAllDefinitions()).toHaveLength(1);
    });

-    it('should log overwrites when in debug mode', () => {
+    it('should log overwrites when in debug mode', async () => {
      const debugConfig = makeFakeConfig({ debugMode: true });
      const debugRegistry = new TestableAgentRegistry(debugConfig);
      const debugLogSpy = vi
        .spyOn(debugLogger, 'log')
        .mockImplementation(() => {});

-      debugRegistry.testRegisterAgent(MOCK_AGENT_V1);
-      debugRegistry.testRegisterAgent(MOCK_AGENT_V2);
+      await debugRegistry.testRegisterAgent(MOCK_AGENT_V1);
+      await debugRegistry.testRegisterAgent(MOCK_AGENT_V2);

      expect(debugLogSpy).toHaveBeenCalledWith(
        `[AgentRegistry] Overriding agent 'MockAgent'`,
      );
    });

-    it('should not log overwrites when not in debug mode', () => {
+    it('should not log overwrites when not in debug mode', async () => {
      const debugLogSpy = vi
        .spyOn(debugLogger, 'log')
        .mockImplementation(() => {});

-      registry.testRegisterAgent(MOCK_AGENT_V1);
-      registry.testRegisterAgent(MOCK_AGENT_V2);
+      await registry.testRegisterAgent(MOCK_AGENT_V1);
+      await registry.testRegisterAgent(MOCK_AGENT_V2);

      expect(debugLogSpy).not.toHaveBeenCalledWith(
        `[AgentRegistry] Overriding agent 'MockAgent'`,
@@ -373,12 +389,10 @@ describe('AgentRegistry', () => {

    it('should handle bulk registrations correctly', async () => {
      const promises = Array.from({ length: 100 }, (_, i) =>
-        Promise.resolve(
-          registry.testRegisterAgent({
-            ...MOCK_AGENT_V1,
-            name: `Agent${i}`,
-          }),
-        ),
+        registry.testRegisterAgent({
+          ...MOCK_AGENT_V1,
+          name: `Agent${i}`,
+        }),
      );

      await Promise.all(promises);
@@ -387,7 +401,7 @@ describe('AgentRegistry', () => {
  });

  describe('inheritance and refresh', () => {
-    it('should resolve "inherit" to the current model from configuration', () => {
+    it('should resolve "inherit" to the current model from configuration', async () => {
      const config = makeFakeConfig({ model: 'current-model' });
      const registry = new TestableAgentRegistry(config);

@@ -396,7 +410,7 @@ describe('AgentRegistry', () => {
        modelConfig: { ...MOCK_AGENT_V1.modelConfig, model: 'inherit' },
      };

-      registry.testRegisterAgent(agent);
+      await registry.testRegisterAgent(agent);

      const resolved = config.modelConfigService.getResolvedConfig({
        model: getModelConfigAlias(agent),
@@ -415,7 +429,7 @@ describe('AgentRegistry', () => {
        modelConfig: { ...MOCK_AGENT_V1.modelConfig, model: 'inherit' },
      };

-      registry.testRegisterAgent(agent);
+      await registry.testRegisterAgent(agent);

      // Verify initial state
      let resolved = config.modelConfigService.getResolvedConfig({
@@ -429,6 +443,17 @@ describe('AgentRegistry', () => {
        model: 'new-model',
      });

+      // Since the listener is async but not awaited by emit, we should manually
+      // trigger refresh or wait.
+      await vi.waitFor(() => {
+        const resolved = config.modelConfigService.getResolvedConfig({
+          model: getModelConfigAlias(agent),
+        });
+        if (resolved.model !== 'new-model') {
+          throw new Error('Model not updated yet');
+        }
+      });
+
      // Verify refreshed state
      resolved = config.modelConfigService.getResolvedConfig({
        model: getModelConfigAlias(agent),
@@ -443,9 +468,9 @@ describe('AgentRegistry', () => {
      name: 'AnotherAgent',
    };

-    beforeEach(() => {
-      registry.testRegisterAgent(MOCK_AGENT_V1);
-      registry.testRegisterAgent(ANOTHER_AGENT);
+    beforeEach(async () => {
+      await registry.testRegisterAgent(MOCK_AGENT_V1);
+      await registry.testRegisterAgent(ANOTHER_AGENT);
    });

    it('getDefinition should return the correct definition', () => {
@@ -472,9 +497,9 @@ describe('AgentRegistry', () => {
      );
    });

-    it('should return formatted list of agents when agents are available', () => {
-      registry.testRegisterAgent(MOCK_AGENT_V1);
-      registry.testRegisterAgent({
+    it('should return formatted list of agents when agents are available', async () => {
+      await registry.testRegisterAgent(MOCK_AGENT_V1);
+      await registry.testRegisterAgent({
        ...MOCK_AGENT_V2,
        name: 'AnotherAgent',
        description: 'Another agent description',
--- a/packages/core/src/agents/registry.ts
+++ b/packages/core/src/agents/registry.ts
@@ -11,6 +11,7 @@ import type { AgentDefinition } from './types.js';
 import { loadAgentsFromDirectory } from './toml-loader.js';
 import { CodebaseInvestigatorAgent } from './codebase-investigator.js';
 import { IntrospectionAgent } from './introspection-agent.js';
+import { A2AClientManager } from './a2a-client-manager.js';
 import { type z } from 'zod';
 import { debugLogger } from '../utils/debugLogger.js';
 import {
@@ -47,7 +48,12 @@ export class AgentRegistry {
    this.loadBuiltInAgents();

    coreEvents.on(CoreEvent.ModelChanged, () => {
-      this.refreshAgents();
+      this.refreshAgents().catch((e) => {
+        debugLogger.error(
+          '[AgentRegistry] Failed to refresh agents on model change:',
+          e,
+        );
+      });
    });

    if (!this.config.isAgentsEnabled()) {
@@ -63,9 +69,9 @@ export class AgentRegistry {
      );
      coreEvents.emitFeedback('error', `Agent loading error: ${error.message}`);
    }
-    for (const agent of userAgents.agents) {
-      this.registerAgent(agent);
-    }
+    await Promise.allSettled(
+      userAgents.agents.map((agent) => this.registerAgent(agent)),
+    );

    // Load project-level agents: .gemini/agents/ (relative to Project Root)
    const folderTrustEnabled = this.config.getFolderTrust();
@@ -80,9 +86,9 @@ export class AgentRegistry {
          `Agent loading error: ${error.message}`,
        );
      }
-      for (const agent of projectAgents.agents) {
-        this.registerAgent(agent);
-      }
+      await Promise.allSettled(
+        projectAgents.agents.map((agent) => this.registerAgent(agent)),
+      );
    } else {
      coreEvents.emitFeedback(
        'info',
@@ -135,20 +141,22 @@ export class AgentRegistry {
            CodebaseInvestigatorAgent.runConfig.max_turns,
        },
      };
-      this.registerAgent(agentDef);
+      this.registerLocalAgent(agentDef);
    }

    // Register the introspection agent if it's explicitly enabled.
    if (introspectionSettings.enabled) {
-      this.registerAgent(IntrospectionAgent);
+      this.registerLocalAgent(IntrospectionAgent);
    }
  }

-  private refreshAgents(): void {
+  private async refreshAgents(): Promise<void> {
    this.loadBuiltInAgents();
-    for (const agent of this.agents.values()) {
-      this.registerAgent(agent);
-    }
+    await Promise.allSettled(
+      Array.from(this.agents.values()).map((agent) =>
+        this.registerAgent(agent),
+      ),
+    );
  }

  /**
@@ -156,9 +164,26 @@ export class AgentRegistry {
   * it will be overwritten, respecting the precedence established by the
   * initialization order.
   */
-  protected registerAgent<TOutput extends z.ZodTypeAny>(
+  protected async registerAgent<TOutput extends z.ZodTypeAny>(
+    definition: AgentDefinition<TOutput>,
+  ): Promise<void> {
+    if (definition.kind === 'local') {
+      this.registerLocalAgent(definition);
+    } else if (definition.kind === 'remote') {
+      await this.registerRemoteAgent(definition);
+    }
+  }
+
+  /**
+   * Registers a local agent definition synchronously.
+   */
+  protected registerLocalAgent<TOutput extends z.ZodTypeAny>(
    definition: AgentDefinition<TOutput>,
  ): void {
+    if (definition.kind !== 'local') {
+      return;
+    }
+
    // Basic validation
    if (!definition.name || !definition.description) {
      debugLogger.warn(
@@ -175,37 +200,79 @@ export class AgentRegistry {

    // Register model config.
    // TODO(12916): Migrate sub-agents where possible to static configs.
-    if (definition.kind === 'local') {
-      const modelConfig = definition.modelConfig;
-      let model = modelConfig.model;
-      if (model === 'inherit') {
-        model = this.config.getModel();
-      }
+    const modelConfig = definition.modelConfig;
+    let model = modelConfig.model;
+    if (model === 'inherit') {
+      model = this.config.getModel();
+    }

-      const runtimeAlias: ModelConfigAlias = {
-        modelConfig: {
-          model,
-          generateContentConfig: {
-            temperature: modelConfig.temp,
-            topP: modelConfig.top_p,
-            thinkingConfig: {
-              includeThoughts: true,
-              thinkingBudget: modelConfig.thinkingBudget ?? -1,
-            },
+    const runtimeAlias: ModelConfigAlias = {
+      modelConfig: {
+        model,
+        generateContentConfig: {
+          temperature: modelConfig.temp,
+          topP: modelConfig.top_p,
+          thinkingConfig: {
+            includeThoughts: true,
+            thinkingBudget: modelConfig.thinkingBudget ?? -1,
          },
        },
-      };
+      },
+    };

-      this.config.modelConfigService.registerRuntimeModelConfig(
-        getModelConfigAlias(definition),
-        runtimeAlias,
+    this.config.modelConfigService.registerRuntimeModelConfig(
+      getModelConfigAlias(definition),
+      runtimeAlias,
+    );
+  }
+
+  /**
+   * Registers a remote agent definition asynchronously.
+   */
+  protected async registerRemoteAgent<TOutput extends z.ZodTypeAny>(
+    definition: AgentDefinition<TOutput>,
+  ): Promise<void> {
+    if (definition.kind !== 'remote') {
+      return;
+    }
+
+    // Basic validation
+    if (!definition.name || !definition.description) {
+      debugLogger.warn(
+        `[AgentRegistry] Skipping invalid agent definition. Missing name or description.`,
      );
+      return;
+    }
+
+    if (this.agents.has(definition.name) && this.config.getDebugMode()) {
+      debugLogger.log(`[AgentRegistry] Overriding agent '${definition.name}'`);
    }

    // Log remote A2A agent registration for visibility.
-    if (definition.kind === 'remote' && this.config.getDebugMode()) {
-      debugLogger.log(
-        `[AgentRegistry] Registered remote agent '${definition.name}' with card: ${definition.agentCardUrl}`,
+    try {
+      const clientManager = A2AClientManager.getInstance();
+      const agentCard = await clientManager.loadAgent(
+        definition.name,
+        definition.agentCardUrl,
+      );
+      if (agentCard.skills && agentCard.skills.length > 0) {
+        definition.description = agentCard.skills
+          .map(
+            (skill: { name: string; description: string }) =>
+              `${skill.name}: ${skill.description}`,
+          )
+          .join('\n');
+      }
+      if (this.config.getDebugMode()) {
+        debugLogger.log(
+          `[AgentRegistry] Registered remote agent '${definition.name}' with card: ${definition.agentCardUrl}`,
+        );
+      }
+      this.agents.set(definition.name, definition);
+    } catch (e) {
+      debugLogger.warn(
+        `[AgentRegistry] Error loading A2A agent "${definition.name}":`,
+        e,
      );
    }
  }