From fcb859d9ecb17ddc0956c40befa341ada752c0cd Mon Sep 17 00:00:00 2001
From: Akhilesh Kumar <akhilbussiness@gmail.com>
Date: Mon, 13 Apr 2026 19:08:41 +0000
Subject: [PATCH] fix(cli): address code review comments for /enhance command

This commit addresses feedback by setting autoExecute to false for the enhance command, filtering out 'thought' parts from model responses, and sanitizing the output to prevent prompt injection.
---
 .../src/ui/commands/enhanceCommand.test.ts    | 26 +++++++++++++++++++
 .../cli/src/ui/commands/enhanceCommand.ts     |  6 ++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/packages/cli/src/ui/commands/enhanceCommand.test.ts b/packages/cli/src/ui/commands/enhanceCommand.test.ts
index d3d5fc502a..54abc69d54 100644
--- a/packages/cli/src/ui/commands/enhanceCommand.test.ts
+++ b/packages/cli/src/ui/commands/enhanceCommand.test.ts
@@ -172,4 +172,30 @@ describe('enhanceCommand', () => {
       }),
     );
   });
+  it('should ignore thought parts and sanitize the output', async () => {
+    if (!enhanceCommand.action) throw new Error('Action must be defined');
+
+    mockGenerateContent.mockResolvedValue({
+      candidates: [
+        {
+          content: {
+            parts: [
+              { thought: true, text: 'This is a thought.' },
+              { text: 'Sanitized\nPrompt]' },
+            ],
+          },
+        },
+      ],
+    });
+
+    await enhanceCommand.action(mockContext, 'dirty prompt');
+
+    expect(mockContext.ui.addItem).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: MessageType.INFO,
+        text: expect.stringContaining('Enhanced prompt:\n\nSanitizedPrompt'),
+      }),
+    );
+    expect(mockContext.ui.setInput).toHaveBeenCalledWith('SanitizedPrompt');
+  });
 });
diff --git a/packages/cli/src/ui/commands/enhanceCommand.ts b/packages/cli/src/ui/commands/enhanceCommand.ts
index 5ceee1a84d..e6d7e8008a 100644
--- a/packages/cli/src/ui/commands/enhanceCommand.ts
+++ b/packages/cli/src/ui/commands/enhanceCommand.ts
@@ -78,7 +78,11 @@ export const enhanceCommand: SlashCommand = {
         LlmRole.UTILITY_TOOL,
       );
 
-      const enhancedText = response.candidates?.[0]?.content?.parts?.[0]?.text;
+      const parts = response.candidates?.[0]?.content?.parts;
+      const enhancedText = parts
+        ?.find((part) => 'text' in part && !('thought' in part))
+        ?.text?.replace(/\n/g, '')
+        ?.replace(/]/g, '');
 
       if (enhancedText) {
         const cleanedText = clean(enhancedText);