From 18dd399cb571e47178dda7fc811d9f4a1867991c Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Thu, 8 Jan 2026 19:51:18 -0500 Subject: [PATCH] Support @ suggestions for subagenets (#16201) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/cli/src/ui/commands/types.ts | 1 + .../src/ui/components/SuggestionsDisplay.tsx | 16 +- .../cli/src/ui/hooks/atCommandProcessor.ts | 22 +- .../hooks/atCommandProcessor_agents.test.ts | 241 ++++++++++++++++++ packages/cli/src/ui/hooks/useAtCompletion.ts | 40 +++ .../ui/hooks/useAtCompletion_agents.test.ts | 117 +++++++++ packages/core/src/index.ts | 3 + 7 files changed, 435 insertions(+), 5 deletions(-) create mode 100644 packages/cli/src/ui/hooks/atCommandProcessor_agents.test.ts create mode 100644 packages/cli/src/ui/hooks/useAtCompletion_agents.test.ts diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index 2165ab377a..a34ff960bb 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -163,6 +163,7 @@ export enum CommandKind { BUILT_IN = 'built-in', FILE = 'file', MCP_PROMPT = 'mcp-prompt', + AGENT = 'agent', } // The standardized contract for any command in the system. diff --git a/packages/cli/src/ui/components/SuggestionsDisplay.tsx b/packages/cli/src/ui/components/SuggestionsDisplay.tsx index 1a8c7cab81..96eb554076 100644 --- a/packages/cli/src/ui/components/SuggestionsDisplay.tsx +++ b/packages/cli/src/ui/components/SuggestionsDisplay.tsx @@ -60,8 +60,13 @@ export function SuggestionsDisplay({ ); const visibleSuggestions = suggestions.slice(startIndex, endIndex); + const COMMAND_KIND_SUFFIX: Partial> = { + [CommandKind.MCP_PROMPT]: ' [MCP]', + [CommandKind.AGENT]: ' [Agent]', + }; + const getFullLabel = (s: Suggestion) => - s.label + (s.commandKind === CommandKind.MCP_PROMPT ? ' [MCP]' : ''); + s.label + (s.commandKind ? (COMMAND_KIND_SUFFIX[s.commandKind] ?? '') : ''); const maxLabelLength = Math.max( ...suggestions.map((s) => getFullLabel(s).length), @@ -98,9 +103,12 @@ export function SuggestionsDisplay({ > {labelElement} - {suggestion.commandKind === CommandKind.MCP_PROMPT && ( - [MCP] - )} + {suggestion.commandKind && + COMMAND_KIND_SUFFIX[suggestion.commandKind] && ( + + {COMMAND_KIND_SUFFIX[suggestion.commandKind]} + + )} diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts index 10196a3545..f545c3e103 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts @@ -155,6 +155,7 @@ export async function handleAtCommand({ const pathSpecsToRead: string[] = []; const resourceAttachments: DiscoveredMCPResource[] = []; const atPathToResolvedSpecMap = new Map(); + const agentsFound: string[] = []; const fileLabelsForDisplay: string[] = []; const absoluteToRelativePathMap = new Map(); const ignoredByReason: Record = { @@ -208,6 +209,14 @@ export async function handleAtCommand({ return { processedQuery: null, error: errMsg }; } + // Check if this is an Agent reference + const agentRegistry = config.getAgentRegistry?.(); + if (agentRegistry?.getDefinition(pathName)) { + agentsFound.push(pathName); + atPathToResolvedSpecMap.set(originalAtPath, pathName); + continue; + } + // Check if this is an MCP resource reference (serverName:uri format) const resourceMatch = resourceRegistry.findResourceByUri(pathName); if (resourceMatch) { @@ -420,7 +429,11 @@ export async function handleAtCommand({ } // Fallback for lone "@" or completely invalid @-commands resulting in empty initialQueryText - if (pathSpecsToRead.length === 0 && resourceAttachments.length === 0) { + if ( + pathSpecsToRead.length === 0 && + resourceAttachments.length === 0 && + agentsFound.length === 0 + ) { onDebugMessage('No valid file paths found in @ commands to read.'); if (initialQueryText === '@' && query.trim() === '@') { // If the only thing was a lone @, pass original query (which might have spaces) @@ -435,6 +448,13 @@ export async function handleAtCommand({ const processedQueryParts: PartListUnion = [{ text: initialQueryText }]; + if (agentsFound.length > 0) { + const agentNudge = `\n\nThe user has explicitly selected the following agent(s): ${agentsFound.join( + ', ', + )}. Please use the 'delegate_to_agent' tool to delegate the task to the selected agent(s).\n\n`; + processedQueryParts.push({ text: agentNudge }); + } + const resourcePromises = resourceAttachments.map(async (resource) => { const uri = resource.uri; const client = mcpClientManager?.getClient(resource.serverName); diff --git a/packages/cli/src/ui/hooks/atCommandProcessor_agents.test.ts b/packages/cli/src/ui/hooks/atCommandProcessor_agents.test.ts new file mode 100644 index 0000000000..0364cf94f6 --- /dev/null +++ b/packages/cli/src/ui/hooks/atCommandProcessor_agents.test.ts @@ -0,0 +1,241 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { handleAtCommand } from './atCommandProcessor.js'; +import type { + Config, + AgentDefinition, + MessageBus, +} from '@google/gemini-cli-core'; +import { + FileDiscoveryService, + GlobTool, + ReadManyFilesTool, + StandardFileSystemService, + ToolRegistry, + COMMON_IGNORE_PATTERNS, +} from '@google/gemini-cli-core'; +import * as os from 'node:os'; +import type { UseHistoryManagerReturn } from './useHistoryManager.js'; +import * as fsPromises from 'node:fs/promises'; +import * as path from 'node:path'; + +describe('handleAtCommand with Agents', () => { + let testRootDir: string; + let mockConfig: Config; + + const mockAddItem: UseHistoryManagerReturn['addItem'] = vi.fn(); + const mockOnDebugMessage: (message: string) => void = vi.fn(); + + let abortController: AbortController; + + beforeEach(async () => { + vi.resetAllMocks(); + + testRootDir = await fsPromises.mkdtemp( + path.join(os.tmpdir(), 'agent-test-'), + ); + + abortController = new AbortController(); + + const getToolRegistry = vi.fn(); + const mockMessageBus = { + publish: vi.fn(), + subscribe: vi.fn(), + unsubscribe: vi.fn(), + } as unknown as MessageBus; + + const mockAgentRegistry = { + getDefinition: vi.fn((name: string) => { + if (name === 'CodebaseInvestigator') { + return { + name: 'CodebaseInvestigator', + description: 'Investigates codebase', + kind: 'local', + } as AgentDefinition; + } + return undefined; + }), + }; + + mockConfig = { + getToolRegistry, + getTargetDir: () => testRootDir, + isSandboxed: () => false, + getExcludeTools: vi.fn(), + getFileService: () => new FileDiscoveryService(testRootDir), + getFileFilteringRespectGitIgnore: () => true, + getFileFilteringRespectGeminiIgnore: () => true, + getFileFilteringOptions: () => ({ + respectGitIgnore: true, + respectGeminiIgnore: true, + }), + getFileSystemService: () => new StandardFileSystemService(), + getEnableRecursiveFileSearch: vi.fn(() => true), + getWorkspaceContext: () => ({ + isPathWithinWorkspace: () => true, + getDirectories: () => [testRootDir], + }), + getMcpServers: () => ({}), + getMcpServerCommand: () => undefined, + getPromptRegistry: () => ({ + getPromptsByServer: () => [], + }), + getDebugMode: () => false, + getFileExclusions: () => ({ + getCoreIgnorePatterns: () => COMMON_IGNORE_PATTERNS, + getDefaultExcludePatterns: () => [], + getGlobExcludes: () => [], + buildExcludePatterns: () => [], + getReadManyFilesExcludes: () => [], + }), + getUsageStatisticsEnabled: () => false, + getEnableExtensionReloading: () => false, + getResourceRegistry: () => ({ + findResourceByUri: () => undefined, + getAllResources: () => [], + }), + getMcpClientManager: () => ({ + getClient: () => undefined, + }), + getAgentRegistry: () => mockAgentRegistry, + getMessageBus: () => mockMessageBus, + } as unknown as Config; + + const registry = new ToolRegistry(mockConfig, mockMessageBus); + registry.registerTool(new ReadManyFilesTool(mockConfig, mockMessageBus)); + registry.registerTool(new GlobTool(mockConfig, mockMessageBus)); + getToolRegistry.mockReturnValue(registry); + }); + + afterEach(async () => { + abortController.abort(); + await fsPromises.rm(testRootDir, { recursive: true, force: true }); + }); + + it('should detect agent reference and add nudge message', async () => { + const query = 'Please help me @CodebaseInvestigator'; + + const result = await handleAtCommand({ + query, + config: mockConfig, + addItem: mockAddItem, + onDebugMessage: mockOnDebugMessage, + messageId: 123, + signal: abortController.signal, + }); + + expect(result.processedQuery).toBeDefined(); + const parts = result.processedQuery; + + if (!Array.isArray(parts)) { + throw new Error('processedQuery should be an array'); + } + + // Check if the query text is preserved + const firstPart = parts[0]; + if ( + typeof firstPart === 'object' && + firstPart !== null && + 'text' in firstPart + ) { + expect((firstPart as { text: string }).text).toContain( + 'Please help me @CodebaseInvestigator', + ); + } else { + throw new Error('First part should be a text part'); + } + + // Check if the nudge message is added + const nudgePart = parts.find( + (p) => + typeof p === 'object' && + p !== null && + 'text' in p && + (p as { text: string }).text.includes(''), + ); + expect(nudgePart).toBeDefined(); + if (nudgePart && typeof nudgePart === 'object' && 'text' in nudgePart) { + expect((nudgePart as { text: string }).text).toContain( + 'The user has explicitly selected the following agent(s): CodebaseInvestigator', + ); + } + }); + + it('should handle multiple agents', async () => { + // Mock another agent + const mockAgentRegistry = mockConfig.getAgentRegistry() as { + getDefinition: (name: string) => AgentDefinition | undefined; + }; + mockAgentRegistry.getDefinition = vi.fn((name: string) => { + if (name === 'CodebaseInvestigator' || name === 'AnotherAgent') { + return { name, description: 'desc', kind: 'local' } as AgentDefinition; + } + return undefined; + }); + + const query = '@CodebaseInvestigator and @AnotherAgent'; + const result = await handleAtCommand({ + query, + config: mockConfig, + addItem: mockAddItem, + onDebugMessage: mockOnDebugMessage, + messageId: 124, + signal: abortController.signal, + }); + + const parts = result.processedQuery; + if (!Array.isArray(parts)) { + throw new Error('processedQuery should be an array'); + } + + const nudgePart = parts.find( + (p) => + typeof p === 'object' && + p !== null && + 'text' in p && + (p as { text: string }).text.includes(''), + ); + expect(nudgePart).toBeDefined(); + if (nudgePart && typeof nudgePart === 'object' && 'text' in nudgePart) { + expect((nudgePart as { text: string }).text).toContain( + 'CodebaseInvestigator, AnotherAgent', + ); + } + }); + + it('should not treat non-agents as agents', async () => { + const query = '@UnknownAgent'; + // This should fail to resolve and fallback or error depending on file search + // Since it's not a file, handleAtCommand logic for files will run. + // It will likely log debug message about not finding file/glob. + // But critical for this test: it should NOT add the agent nudge. + + const result = await handleAtCommand({ + query, + config: mockConfig, + addItem: mockAddItem, + onDebugMessage: mockOnDebugMessage, + messageId: 125, + signal: abortController.signal, + }); + + const parts = result.processedQuery; + if (!Array.isArray(parts)) { + throw new Error('processedQuery should be an array'); + } + + const nudgePart = parts.find( + (p) => + typeof p === 'object' && + p !== null && + 'text' in p && + (p as { text: string }).text.includes(''), + ); + expect(nudgePart).toBeUndefined(); + }); +}); diff --git a/packages/cli/src/ui/hooks/useAtCompletion.ts b/packages/cli/src/ui/hooks/useAtCompletion.ts index c361089947..dcb6dfa478 100644 --- a/packages/cli/src/ui/hooks/useAtCompletion.ts +++ b/packages/cli/src/ui/hooks/useAtCompletion.ts @@ -9,6 +9,7 @@ import type { Config, FileSearch } from '@google/gemini-cli-core'; import { FileSearchFactory, escapePath } from '@google/gemini-cli-core'; import type { Suggestion } from '../components/SuggestionsDisplay.js'; import { MAX_SUGGESTIONS_TO_SHOW } from '../components/SuggestionsDisplay.js'; +import { CommandKind } from '../commands/types.js'; import { AsyncFzf } from 'fzf'; export enum AtCompletionStatus { @@ -127,6 +128,18 @@ function buildResourceCandidates( return resources; } +function buildAgentCandidates(config?: Config): Suggestion[] { + const registry = config?.getAgentRegistry?.(); + if (!registry) { + return []; + } + return registry.getAllDefinitions().map((def) => ({ + label: def.name, + value: def.name, + commandKind: CommandKind.AGENT, + })); +} + async function searchResourceCandidates( pattern: string, candidates: ResourceSuggestionCandidate[], @@ -153,6 +166,26 @@ async function searchResourceCandidates( ); } +async function searchAgentCandidates( + pattern: string, + candidates: Suggestion[], +): Promise { + if (candidates.length === 0) { + return []; + } + const normalizedPattern = pattern.toLowerCase(); + if (!normalizedPattern) { + return candidates.slice(0, MAX_SUGGESTIONS_TO_SHOW); + } + const fzf = new AsyncFzf(candidates, { + selector: (s: Suggestion) => s.label, + }); + const results = await fzf.find(normalizedPattern, { + limit: MAX_SUGGESTIONS_TO_SHOW, + }); + return results.map((r: { item: Suggestion }) => r.item); +} + export function useAtCompletion(props: UseAtCompletionProps): void { const { enabled, @@ -283,7 +316,14 @@ export function useAtCompletion(props: UseAtCompletionProps): void { value: suggestion.value.replace(/^@/, ''), })); + const agentCandidates = buildAgentCandidates(config); + const agentSuggestions = await searchAgentCandidates( + state.pattern ?? '', + agentCandidates, + ); + const combinedSuggestions = [ + ...agentSuggestions, ...fileSuggestions, ...resourceSuggestions, ]; diff --git a/packages/cli/src/ui/hooks/useAtCompletion_agents.test.ts b/packages/cli/src/ui/hooks/useAtCompletion_agents.test.ts new file mode 100644 index 0000000000..53b416ce6a --- /dev/null +++ b/packages/cli/src/ui/hooks/useAtCompletion_agents.test.ts @@ -0,0 +1,117 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'; +import { useState } from 'react'; +import { renderHook } from '../../test-utils/render.js'; +import { waitFor } from '../../test-utils/async.js'; +import { useAtCompletion } from './useAtCompletion.js'; +import type { Config, AgentDefinition } from '@google/gemini-cli-core'; +import { createTmpDir, cleanupTmpDir } from '@google/gemini-cli-test-utils'; +import type { Suggestion } from '../components/SuggestionsDisplay.js'; +import { CommandKind } from '../commands/types.js'; + +// Test harness to capture the state from the hook's callbacks. +function useTestHarnessForAtCompletion( + enabled: boolean, + pattern: string, + config: Config | undefined, + cwd: string, +) { + const [suggestions, setSuggestions] = useState([]); + const [isLoadingSuggestions, setIsLoadingSuggestions] = useState(false); + + useAtCompletion({ + enabled, + pattern, + config, + cwd, + setSuggestions, + setIsLoadingSuggestions, + }); + + return { suggestions, isLoadingSuggestions }; +} + +describe('useAtCompletion with Agents', () => { + let testRootDir: string; + let mockConfig: Config; + + beforeEach(() => { + const mockAgentRegistry = { + getAllDefinitions: vi.fn(() => [ + { + name: 'CodebaseInvestigator', + description: 'Investigates codebase', + kind: 'local', + } as AgentDefinition, + { + name: 'OtherAgent', + description: 'Another agent', + kind: 'local', + } as AgentDefinition, + ]), + }; + + mockConfig = { + getFileFilteringOptions: vi.fn(() => ({ + respectGitIgnore: true, + respectGeminiIgnore: true, + })), + getEnableRecursiveFileSearch: () => true, + getFileFilteringDisableFuzzySearch: () => false, + getResourceRegistry: vi.fn().mockReturnValue({ + getAllResources: () => [], + }), + getAgentRegistry: () => mockAgentRegistry, + } as unknown as Config; + vi.clearAllMocks(); + }); + + afterEach(async () => { + if (testRootDir) { + await cleanupTmpDir(testRootDir); + } + vi.restoreAllMocks(); + }); + + it('should include agent suggestions', async () => { + testRootDir = await createTmpDir({}); + + const { result } = renderHook(() => + useTestHarnessForAtCompletion(true, '', mockConfig, testRootDir), + ); + + await waitFor(() => { + expect(result.current.suggestions.length).toBeGreaterThan(0); + }); + + const agentSuggestion = result.current.suggestions.find( + (s) => s.value === 'CodebaseInvestigator', + ); + expect(agentSuggestion).toBeDefined(); + expect(agentSuggestion?.commandKind).toBe(CommandKind.AGENT); + }); + + it('should filter agent suggestions', async () => { + testRootDir = await createTmpDir({}); + + const { result } = renderHook(() => + useTestHarnessForAtCompletion(true, 'Code', mockConfig, testRootDir), + ); + + await waitFor(() => { + expect(result.current.suggestions.length).toBeGreaterThan(0); + }); + + expect(result.current.suggestions.map((s) => s.value)).toContain( + 'CodebaseInvestigator', + ); + expect(result.current.suggestions.map((s) => s.value)).not.toContain( + 'OtherAgent', + ); + }); +}); diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index a56d5ae813..75acd00143 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -167,6 +167,9 @@ export * from './hooks/index.js'; // Export hook types export * from './hooks/types.js'; +// Export agent types +export * from './agents/types.js'; + // Export stdio utils export * from './utils/stdio.js'; export * from './utils/terminal.js';