diff --git a/packages/cli/src/commands/extensions.test.tsx b/packages/cli/src/commands/extensions.test.tsx index d8aae8b359..0630b398ff 100644 --- a/packages/cli/src/commands/extensions.test.tsx +++ b/packages/cli/src/commands/extensions.test.tsx @@ -54,15 +54,33 @@ describe('extensionsCommand', () => { extensionsCommand.builder(mockYargs); expect(mockYargs.middleware).toHaveBeenCalled(); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'install' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'uninstall' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'list' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'update' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'disable' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'enable' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'link' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'new' }); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'validate' }); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'install' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'uninstall' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'list' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'update' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'disable' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'enable' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'link' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'new' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'validate' }), + ); expect(mockYargs.demandCommand).toHaveBeenCalledWith(1, expect.any(String)); expect(mockYargs.version).toHaveBeenCalledWith(false); }); diff --git a/packages/cli/src/commands/extensions.tsx b/packages/cli/src/commands/extensions.tsx index 8079d67256..893af14d11 100644 --- a/packages/cli/src/commands/extensions.tsx +++ b/packages/cli/src/commands/extensions.tsx @@ -16,6 +16,7 @@ import { newCommand } from './extensions/new.js'; import { validateCommand } from './extensions/validate.js'; import { configureCommand } from './extensions/configure.js'; import { initializeOutputListenersAndFlush } from '../gemini.js'; +import { defer } from '../deferred.js'; export const extensionsCommand: CommandModule = { command: 'extensions ', @@ -24,16 +25,16 @@ export const extensionsCommand: CommandModule = { builder: (yargs) => yargs .middleware(() => initializeOutputListenersAndFlush()) - .command(installCommand) - .command(uninstallCommand) - .command(listCommand) - .command(updateCommand) - .command(disableCommand) - .command(enableCommand) - .command(linkCommand) - .command(newCommand) - .command(validateCommand) - .command(configureCommand) + .command(defer(installCommand, 'extensions')) + .command(defer(uninstallCommand, 'extensions')) + .command(defer(listCommand, 'extensions')) + .command(defer(updateCommand, 'extensions')) + .command(defer(disableCommand, 'extensions')) + .command(defer(enableCommand, 'extensions')) + .command(defer(linkCommand, 'extensions')) + .command(defer(newCommand, 'extensions')) + .command(defer(validateCommand, 'extensions')) + .command(defer(configureCommand, 'extensions')) .demandCommand(1, 'You need at least one command before continuing.') .version(false), handler: () => { diff --git a/packages/cli/src/commands/mcp.ts b/packages/cli/src/commands/mcp.ts index f09680dbdd..c14def2199 100644 --- a/packages/cli/src/commands/mcp.ts +++ b/packages/cli/src/commands/mcp.ts @@ -10,6 +10,7 @@ import { addCommand } from './mcp/add.js'; import { removeCommand } from './mcp/remove.js'; import { listCommand } from './mcp/list.js'; import { initializeOutputListenersAndFlush } from '../gemini.js'; +import { defer } from '../deferred.js'; export const mcpCommand: CommandModule = { command: 'mcp', @@ -17,9 +18,9 @@ export const mcpCommand: CommandModule = { builder: (yargs: Argv) => yargs .middleware(() => initializeOutputListenersAndFlush()) - .command(addCommand) - .command(removeCommand) - .command(listCommand) + .command(defer(addCommand, 'mcp')) + .command(defer(removeCommand, 'mcp')) + .command(defer(listCommand, 'mcp')) .demandCommand(1, 'You need at least one command before continuing.') .version(false), handler: () => { diff --git a/packages/cli/src/commands/skills.test.tsx b/packages/cli/src/commands/skills.test.tsx index 5a76ab0d95..e7b9a4eb9d 100644 --- a/packages/cli/src/commands/skills.test.tsx +++ b/packages/cli/src/commands/skills.test.tsx @@ -38,13 +38,19 @@ describe('skillsCommand', () => { skillsCommand.builder(mockYargs); expect(mockYargs.middleware).toHaveBeenCalled(); - expect(mockYargs.command).toHaveBeenCalledWith({ command: 'list' }); - expect(mockYargs.command).toHaveBeenCalledWith({ - command: 'enable ', - }); - expect(mockYargs.command).toHaveBeenCalledWith({ - command: 'disable ', - }); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ command: 'list' }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ + command: 'enable ', + }), + ); + expect(mockYargs.command).toHaveBeenCalledWith( + expect.objectContaining({ + command: 'disable ', + }), + ); expect(mockYargs.demandCommand).toHaveBeenCalledWith(1, expect.any(String)); expect(mockYargs.version).toHaveBeenCalledWith(false); }); diff --git a/packages/cli/src/commands/skills.tsx b/packages/cli/src/commands/skills.tsx index d2f249b049..b8879681cd 100644 --- a/packages/cli/src/commands/skills.tsx +++ b/packages/cli/src/commands/skills.tsx @@ -11,6 +11,7 @@ import { disableCommand } from './skills/disable.js'; import { installCommand } from './skills/install.js'; import { uninstallCommand } from './skills/uninstall.js'; import { initializeOutputListenersAndFlush } from '../gemini.js'; +import { defer } from '../deferred.js'; export const skillsCommand: CommandModule = { command: 'skills ', @@ -19,11 +20,11 @@ export const skillsCommand: CommandModule = { builder: (yargs) => yargs .middleware(() => initializeOutputListenersAndFlush()) - .command(listCommand) - .command(enableCommand) - .command(disableCommand) - .command(installCommand) - .command(uninstallCommand) + .command(defer(listCommand, 'skills')) + .command(defer(enableCommand, 'skills')) + .command(defer(disableCommand, 'skills')) + .command(defer(installCommand, 'skills')) + .command(defer(uninstallCommand, 'skills')) .demandCommand(1, 'You need at least one command before continuing.') .version(false), handler: () => { diff --git a/packages/cli/src/deferred.test.ts b/packages/cli/src/deferred.test.ts new file mode 100644 index 0000000000..4ea5eb791d --- /dev/null +++ b/packages/cli/src/deferred.test.ts @@ -0,0 +1,217 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { + runDeferredCommand, + defer, + setDeferredCommand, + type DeferredCommand, +} from './deferred.js'; +import { ExitCodes } from '@google/gemini-cli-core'; +import type { ArgumentsCamelCase, CommandModule } from 'yargs'; +import type { MergedSettings } from './config/settings.js'; +import type { MockInstance } from 'vitest'; + +const { mockRunExitCleanup, mockDebugLogger } = vi.hoisted(() => ({ + mockRunExitCleanup: vi.fn(), + mockDebugLogger: { + log: vi.fn(), + error: vi.fn(), + }, +})); + +vi.mock('@google/gemini-cli-core', async () => { + const actual = await vi.importActual('@google/gemini-cli-core'); + return { + ...actual, + debugLogger: mockDebugLogger, + }; +}); + +vi.mock('./utils/cleanup.js', () => ({ + runExitCleanup: mockRunExitCleanup, +})); + +let mockExit: MockInstance; + +describe('deferred', () => { + beforeEach(() => { + vi.clearAllMocks(); + mockExit = vi + .spyOn(process, 'exit') + .mockImplementation(() => undefined as never); + setDeferredCommand(undefined as unknown as DeferredCommand); // Reset deferred command + }); + + const createMockSettings = (adminSettings: unknown = {}): MergedSettings => + ({ + admin: adminSettings, + }) as unknown as MergedSettings; + + describe('runDeferredCommand', () => { + it('should do nothing if no deferred command is set', async () => { + await runDeferredCommand(createMockSettings()); + expect(mockDebugLogger.log).not.toHaveBeenCalled(); + expect(mockDebugLogger.error).not.toHaveBeenCalled(); + expect(mockExit).not.toHaveBeenCalled(); + }); + + it('should execute the deferred command if enabled', async () => { + const mockHandler = vi.fn(); + setDeferredCommand({ + handler: mockHandler, + argv: { _: [], $0: 'gemini' } as ArgumentsCamelCase, + commandName: 'mcp', + }); + + const settings = createMockSettings({ mcp: { enabled: true } }); + await runDeferredCommand(settings); + expect(mockHandler).toHaveBeenCalled(); + expect(mockRunExitCleanup).toHaveBeenCalled(); + expect(mockExit).toHaveBeenCalledWith(ExitCodes.SUCCESS); + }); + + it('should exit with FATAL_CONFIG_ERROR if MCP is disabled', async () => { + setDeferredCommand({ + handler: vi.fn(), + argv: {} as ArgumentsCamelCase, + commandName: 'mcp', + }); + + const settings = createMockSettings({ mcp: { enabled: false } }); + await runDeferredCommand(settings); + + expect(mockDebugLogger.error).toHaveBeenCalledWith( + 'Error: MCP is disabled by your admin.', + ); + expect(mockRunExitCleanup).toHaveBeenCalled(); + expect(mockExit).toHaveBeenCalledWith(ExitCodes.FATAL_CONFIG_ERROR); + }); + + it('should exit with FATAL_CONFIG_ERROR if extensions are disabled', async () => { + setDeferredCommand({ + handler: vi.fn(), + argv: {} as ArgumentsCamelCase, + commandName: 'extensions', + }); + + const settings = createMockSettings({ extensions: { enabled: false } }); + await runDeferredCommand(settings); + + expect(mockDebugLogger.error).toHaveBeenCalledWith( + 'Error: Extensions are disabled by your admin.', + ); + expect(mockRunExitCleanup).toHaveBeenCalled(); + expect(mockExit).toHaveBeenCalledWith(ExitCodes.FATAL_CONFIG_ERROR); + }); + + it('should exit with FATAL_CONFIG_ERROR if skills are disabled', async () => { + setDeferredCommand({ + handler: vi.fn(), + argv: {} as ArgumentsCamelCase, + commandName: 'skills', + }); + + const settings = createMockSettings({ skills: { enabled: false } }); + await runDeferredCommand(settings); + + expect(mockDebugLogger.error).toHaveBeenCalledWith( + 'Error: Agent skills are disabled by your admin.', + ); + expect(mockRunExitCleanup).toHaveBeenCalled(); + expect(mockExit).toHaveBeenCalledWith(ExitCodes.FATAL_CONFIG_ERROR); + }); + + it('should execute if admin settings are undefined (default implicit enable)', async () => { + const mockHandler = vi.fn(); + setDeferredCommand({ + handler: mockHandler, + argv: {} as ArgumentsCamelCase, + commandName: 'mcp', + }); + + const settings = createMockSettings({}); // No admin settings + await runDeferredCommand(settings); + + expect(mockHandler).toHaveBeenCalled(); + expect(mockExit).toHaveBeenCalledWith(ExitCodes.SUCCESS); + }); + }); + + describe('defer', () => { + it('should wrap a command module and defer execution', async () => { + const originalHandler = vi.fn(); + const commandModule: CommandModule = { + command: 'test', + describe: 'test command', + handler: originalHandler, + }; + + const deferredModule = defer(commandModule); + expect(deferredModule.command).toBe(commandModule.command); + + // Execute the wrapper handler + const argv = { _: [], $0: 'gemini' } as ArgumentsCamelCase; + await deferredModule.handler(argv); + + // Should check that it set the deferred command, but didn't run original handler yet + expect(originalHandler).not.toHaveBeenCalled(); + + // Now manually run it to verify it captured correctly + await runDeferredCommand(createMockSettings()); + expect(originalHandler).toHaveBeenCalledWith(argv); + expect(mockExit).toHaveBeenCalledWith(ExitCodes.SUCCESS); + }); + + it('should use parentCommandName if provided', async () => { + const commandModule: CommandModule = { + command: 'subcommand', + describe: 'sub command', + handler: vi.fn(), + }; + + const deferredModule = defer(commandModule, 'parent'); + await deferredModule.handler({} as ArgumentsCamelCase); + + const deferredMcp = defer(commandModule, 'mcp'); + await deferredMcp.handler({} as ArgumentsCamelCase); + + const mcpSettings = createMockSettings({ mcp: { enabled: false } }); + await runDeferredCommand(mcpSettings); + + expect(mockDebugLogger.error).toHaveBeenCalledWith( + 'Error: MCP is disabled by your admin.', + ); + }); + + it('should fallback to unknown if no parentCommandName is provided', async () => { + const mockHandler = vi.fn(); + const commandModule: CommandModule = { + command: ['foo', 'infoo'], + describe: 'foo command', + handler: mockHandler, + }; + + const deferredModule = defer(commandModule); + await deferredModule.handler({} as ArgumentsCamelCase); + + // Verify it runs even if all known commands are disabled, + // confirming it didn't capture 'mcp', 'extensions', or 'skills' + // and defaulted to 'unknown' (or something else safe). + const settings = createMockSettings({ + mcp: { enabled: false }, + extensions: { enabled: false }, + skills: { enabled: false }, + }); + + await runDeferredCommand(settings); + + expect(mockHandler).toHaveBeenCalled(); + expect(mockExit).toHaveBeenCalledWith(ExitCodes.SUCCESS); + }); + }); +}); diff --git a/packages/cli/src/deferred.ts b/packages/cli/src/deferred.ts new file mode 100644 index 0000000000..73fac6d1ce --- /dev/null +++ b/packages/cli/src/deferred.ts @@ -0,0 +1,78 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ +import type { ArgumentsCamelCase, CommandModule } from 'yargs'; +import { debugLogger, ExitCodes } from '@google/gemini-cli-core'; +import { runExitCleanup } from './utils/cleanup.js'; +import type { MergedSettings } from './config/settings.js'; +import process from 'node:process'; + +export interface DeferredCommand { + handler: (argv: ArgumentsCamelCase) => void | Promise; + argv: ArgumentsCamelCase; + commandName: string; +} + +let deferredCommand: DeferredCommand | undefined; + +export function setDeferredCommand(command: DeferredCommand) { + deferredCommand = command; +} + +export async function runDeferredCommand(settings: MergedSettings) { + if (!deferredCommand) { + return; + } + + const adminSettings = settings.admin; + const commandName = deferredCommand.commandName; + + if (commandName === 'mcp' && adminSettings?.mcp?.enabled === false) { + debugLogger.error('Error: MCP is disabled by your admin.'); + await runExitCleanup(); + process.exit(ExitCodes.FATAL_CONFIG_ERROR); + } + + if ( + commandName === 'extensions' && + adminSettings?.extensions?.enabled === false + ) { + debugLogger.error('Error: Extensions are disabled by your admin.'); + await runExitCleanup(); + process.exit(ExitCodes.FATAL_CONFIG_ERROR); + } + + if (commandName === 'skills' && adminSettings?.skills?.enabled === false) { + debugLogger.error('Error: Agent skills are disabled by your admin.'); + await runExitCleanup(); + process.exit(ExitCodes.FATAL_CONFIG_ERROR); + } + + await deferredCommand.handler(deferredCommand.argv); + await runExitCleanup(); + process.exit(ExitCodes.SUCCESS); +} + +/** + * Wraps a command's handler to defer its execution. + * It stores the handler and arguments in a singleton `deferredCommand` variable. + */ +export function defer( + commandModule: CommandModule, + parentCommandName?: string, +): CommandModule { + return { + ...commandModule, + handler: (argv: ArgumentsCamelCase) => { + setDeferredCommand({ + handler: commandModule.handler as ( + argv: ArgumentsCamelCase, + ) => void | Promise, + argv: argv as unknown as ArgumentsCamelCase, + commandName: parentCommandName || 'unknown', + }); + }, + }; +} diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index ee35edae63..733444fdba 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -96,6 +96,7 @@ import { isAlternateBufferEnabled } from './ui/hooks/useAlternateBuffer.js'; import { setupTerminalAndTheme } from './utils/terminalTheme.js'; import { profiler } from './ui/components/DebugProfiler.js'; +import { runDeferredCommand } from './deferred.js'; const SLOW_RENDER_MS = 200; @@ -410,6 +411,9 @@ export async function main() { settings.setRemoteAdminSettings(remoteAdminSettings); } + // Run deferred command now that we have admin settings. + await runDeferredCommand(settings.merged); + // hop into sandbox if we are outside and sandboxing is enabled if (!process.env['SANDBOX']) { const memoryArgs = settings.merged.advanced.autoConfigureMemory