Skip to content
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions .gemini/settings.json
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
{
"experimental": {
"plan": true,
"extensionReloading": true,
"modelSteering": true,
"memoryManager": true
},
"general": {
"devtools": true
"devtools": true,
"plan": {
"enabled": true
}
},
"agents": {
"overrides": {}
}
}
}
2 changes: 2 additions & 0 deletions packages/cli/src/services/BuiltinCommandLoader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import { corgiCommand } from '../ui/commands/corgiCommand.js';
import { docsCommand } from '../ui/commands/docsCommand.js';
import { directoryCommand } from '../ui/commands/directoryCommand.js';
import { editorCommand } from '../ui/commands/editorCommand.js';
import { enhanceCommand } from '../ui/commands/enhanceCommand.js';
import { extensionsCommand } from '../ui/commands/extensionsCommand.js';
import { footerCommand } from '../ui/commands/footerCommand.js';
import { helpCommand } from '../ui/commands/helpCommand.js';
Expand Down Expand Up @@ -133,6 +134,7 @@ export class BuiltinCommandLoader implements ICommandLoader {
docsCommand,
directoryCommand,
editorCommand,
enhanceCommand,
...(this.config?.getExtensionsEnabled() === false
? [
{
Expand Down
1 change: 1 addition & 0 deletions packages/cli/src/test-utils/mockCommandContext.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ export const createMockCommandContext = (
pendingItem: null,
setPendingItem: vi.fn(),
loadHistory: vi.fn(),
setInput: vi.fn(),
toggleCorgiMode: vi.fn(),
toggleShortcutsHelp: vi.fn(),
toggleVimEnabled: vi.fn(),
Expand Down
175 changes: 175 additions & 0 deletions packages/cli/src/ui/commands/enhanceCommand.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,175 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/

import { describe, it, expect, beforeEach, vi, type Mock } from 'vitest';
import { enhanceCommand } from './enhanceCommand.js';
import { type CommandContext } from './types.js';
import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
import { MessageType } from '../types.js';
import { LlmRole } from '@google/gemini-cli-core';

describe('enhanceCommand', () => {
let mockContext: CommandContext;
let mockGenerateContent: Mock;

beforeEach(() => {
mockGenerateContent = vi.fn();
mockContext = createMockCommandContext({
services: {
agentContext: {
promptId: 'test-prompt-id',
geminiClient: {
getHistory: vi.fn().mockReturnValue([
{ role: 'user', parts: [{ text: 'previous user msg' }] },
{ role: 'model', parts: [{ text: 'previous model msg' }] },
]),
},
config: {
getModel: vi.fn().mockReturnValue('test-model'),
getContentGenerator: vi.fn().mockReturnValue({
generateContent: mockGenerateContent,
}),
getGemini31LaunchedSync: vi.fn().mockReturnValue(true),
getHasAccessToPreviewModel: vi.fn().mockReturnValue(true),
},
},
},
ui: {
addItem: vi.fn(),
setDebugMessage: vi.fn(),
},
} as unknown as CommandContext);
});

it('should have the correct name and description', () => {
expect(enhanceCommand.name).toBe('enhance');
expect(enhanceCommand.description).toBe(
'Enhance a prompt with additional context and rephrasing',
);
});

it('should show error if no prompt is provided', async () => {
if (!enhanceCommand.action) throw new Error('Action must be defined');

await enhanceCommand.action(mockContext, '');

expect(mockContext.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.ERROR,
text: expect.stringContaining('Please provide a prompt'),
}),
);
});

it('should call generateContent with correct parameters and show enhanced prompt', async () => {
if (!enhanceCommand.action) throw new Error('Action must be defined');

mockGenerateContent.mockResolvedValue({
candidates: [
{
content: {
parts: [{ text: 'Enhanced: do something' }],
},
},
],
});

await enhanceCommand.action(mockContext, 'do something');

expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
model: 'test-model',
contents: [
{ role: 'user', parts: [{ text: 'previous user msg' }] },
{ role: 'model', parts: [{ text: 'previous model msg' }] },
{ role: 'user', parts: [{ text: 'do something' }] },
],
config: {
systemInstruction: {
role: 'system',
parts: [
{
text: expect.stringContaining(
"Generate an enhanced version of the user's prompt",
),
},
],
},
},
}),
'test-prompt-id',
LlmRole.UTILITY_TOOL,
);

expect(mockContext.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.INFO,
text: expect.stringContaining(
'Enhanced prompt:\n\nEnhanced: do something',
),
}),
);
expect(mockContext.ui.setInput).toHaveBeenCalledWith(
'Enhanced: do something',
);
});

it('should clean the response from markdown and quotes', async () => {
if (!enhanceCommand.action) throw new Error('Action must be defined');

mockGenerateContent.mockResolvedValue({
candidates: [
{
content: {
parts: [{ text: '```markdown\n"Clean me"\n```' }],
},
},
],
});

await enhanceCommand.action(mockContext, 'dirty prompt');

expect(mockContext.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.INFO,
text: expect.stringContaining('Enhanced prompt:\n\nClean me'),
}),
);
expect(mockContext.ui.setInput).toHaveBeenCalledWith('Clean me');
});

it('should handle API errors gracefully', async () => {
if (!enhanceCommand.action) throw new Error('Action must be defined');

mockGenerateContent.mockRejectedValue(new Error('API Error'));

await enhanceCommand.action(mockContext, 'test prompt');

expect(mockContext.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.ERROR,
text: expect.stringContaining('Failed to enhance prompt: API Error'),
}),
);
});

it('should handle empty response from model', async () => {
if (!enhanceCommand.action) throw new Error('Action must be defined');

mockGenerateContent.mockResolvedValue({
candidates: [],
});

await enhanceCommand.action(mockContext, 'test prompt');

expect(mockContext.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.ERROR,
text: expect.stringContaining('Empty response from model'),
}),
);
});
});
104 changes: 104 additions & 0 deletions packages/cli/src/ui/commands/enhanceCommand.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/

import { CommandKind, type SlashCommand } from './types.js';
import { MessageType } from '../types.js';
import { LlmRole, resolveModel } from '@google/gemini-cli-core';
import type { Content } from '@google/genai';

const INSTRUCTION =
"Generate an enhanced version of the user's prompt, using the preceding conversation as context. Reply with ONLY the enhanced prompt - no conversation, explanations, lead-in, bullet points, placeholders, or surrounding quotes.";

function clean(text: string) {
const stripped = text.replace(/^```\w*\n?|```$/g, '').trim();
return stripped.replace(/^(['"])([\s\S]*)\1$/, '$2').trim();
}

export const enhanceCommand: SlashCommand = {
name: 'enhance',
description: 'Enhance a prompt with additional context and rephrasing',
kind: CommandKind.BUILT_IN,
autoExecute: true,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Since the /enhance command requires a prompt argument to function correctly, autoExecute should be set to false. When autoExecute is true, selecting the command from the autocomplete list (e.g., by pressing Enter) will trigger it immediately with empty arguments, resulting in an error message. Setting it to false allows the user to complete the command name and then type the prompt they wish to enhance.

Suggested change
autoExecute: true,
autoExecute: false,

action: async (context, args) => {
const draft = args.trim();
if (!draft) {
context.ui.addItem({
type: MessageType.ERROR,
text: 'Please provide a prompt to enhance. Usage: /enhance <prompt>',
});
return;
}

const agentContext = context.services.agentContext;
if (!agentContext) {
context.ui.addItem({
type: MessageType.ERROR,
text: 'Agent context not available.',
});
return;
}

const config = agentContext.config;
const contentGenerator = config.getContentGenerator();
const promptId = agentContext.promptId;

const model = resolveModel(
config.getModel(),
config.getGemini31LaunchedSync?.() ?? false,
false,
config.getHasAccessToPreviewModel?.() ?? true,
config,
);

context.ui.setDebugMessage('Enhancing prompt...');

try {
const history = agentContext.geminiClient?.getHistory() ?? [];
const contents: Content[] = [
...history,
{ role: 'user', parts: [{ text: draft }] },
];

const response = await contentGenerator.generateContent(
{
model,
contents,
config: {
systemInstruction: {
role: 'system',
parts: [{ text: INSTRUCTION }],
},
},
},
promptId,
LlmRole.UTILITY_TOOL,
);

const enhancedText = response.candidates?.[0]?.content?.parts?.[0]?.text;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-high high

The current implementation assumes the enhanced prompt is always in the first part of the first candidate. Gemini 2.0 models often include a 'thought' part before the actual response text. It is safer to find the first text part that isn't a thought. Additionally, to prevent prompt injection, sanitize the output by removing newlines and context-breaking characters like ']' before it is used in further prompts, as per repository rules.

Suggested change
const enhancedText = response.candidates?.[0]?.content?.parts?.[0]?.text;
const enhancedText = response.candidates?.[0]?.content?.parts?.find((part) => part.text !== undefined && !part.thought)?.text?.replace(/\n/g, '').replace(/]/g, '');
References
  1. Sanitize data from LLM-driven tools before injecting it into a system prompt to prevent prompt injection. At a minimum, remove newlines and context-breaking characters (e.g., ']').


if (enhancedText) {
const cleanedText = clean(enhancedText);
context.ui.addItem({
type: MessageType.INFO,
text: `Enhanced prompt:\n\n${cleanedText}`,
});
context.ui.setInput(cleanedText);
} else {
context.ui.addItem({
type: MessageType.ERROR,
text: 'Failed to enhance prompt: Empty response from model.',
});
}
} catch (error) {
context.ui.addItem({
type: MessageType.ERROR,
text: `Failed to enhance prompt: ${error instanceof Error ? error.message : String(error)}`,
});
} finally {
context.ui.setDebugMessage('');
}
},
};
6 changes: 6 additions & 0 deletions packages/cli/src/ui/commands/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,12 @@ export interface CommandContext {
* @param postLoadInput Optional text to set in the input buffer after loading history.
*/
loadHistory: (history: HistoryItem[], postLoadInput?: string) => void;
/**
* Sets the text in the input buffer.
*
* @param text The text to set.
*/
setInput: (text: string) => void;
/** Toggles a special display mode. */
toggleCorgiMode: () => void;
toggleDebugProfiler: () => void;
Expand Down
1 change: 1 addition & 0 deletions packages/cli/src/ui/hooks/slashCommandProcessor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ export const useSlashCommandProcessor = (
actions.setText(postLoadInput);
}
},
setInput: (text) => actions.setText(text),
setDebugMessage: actions.setDebugMessage,
pendingItem,
setPendingItem,
Expand Down
1 change: 1 addition & 0 deletions packages/cli/src/ui/noninteractive/nonInteractiveUi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ export function createNonInteractiveUI(): CommandContext['ui'] {
clear: () => {},
setDebugMessage: (_message) => {},
loadHistory: (_newHistory) => {},
setInput: (_text) => {},
pendingItem: null,
setPendingItem: (_item) => {},
toggleCorgiMode: () => {},
Expand Down
Loading