From 79900735432f865e9d78108c17068bea4ef8a662 Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Tue, 20 Jan 2026 23:58:37 -0500 Subject: [PATCH] feat(cli): sanitize ANSI escape sequences in non-interactive output (#17172) --- packages/cli/src/config/config.ts | 13 ++ packages/cli/src/gemini.test.tsx | 2 + packages/cli/src/nonInteractiveCli.test.ts | 255 +++++++++++++++++++++ packages/cli/src/nonInteractiveCli.ts | 20 +- packages/core/src/config/config.ts | 14 ++ 5 files changed, 301 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 31cf91412d..65ac1fc4c3 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -83,6 +83,8 @@ export interface CliArgs { outputFormat: string | undefined; fakeResponses: string | undefined; recordResponses: string | undefined; + rawOutput: boolean | undefined; + acceptRawOutputRisk: boolean | undefined; } export async function parseArguments( @@ -248,6 +250,15 @@ export async function parseArguments( type: 'string', description: 'Path to a file to record model responses for testing.', hidden: true, + }) + .option('raw-output', { + type: 'boolean', + description: + 'Disable sanitization of model output (e.g. allow ANSI escape sequences). WARNING: This can be a security risk if the model output is untrusted.', + }) + .option('accept-raw-output-risk', { + type: 'boolean', + description: 'Suppress the security warning when using --raw-output.', }), ) // Register MCP subcommands @@ -759,6 +770,8 @@ export async function loadCliConfig( retryFetchErrors: settings.general?.retryFetchErrors, ptyInfo: ptyInfo?.name, disableLLMCorrection: settings.tools?.disableLLMCorrection, + rawOutput: argv.rawOutput, + acceptRawOutputRisk: argv.acceptRawOutputRisk, modelConfigServiceConfig: settings.modelConfigs, // TODO: loading of hooks based on workspace trust enableHooks: diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 896f89e3c8..ef801786cf 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -490,6 +490,8 @@ describe('gemini.tsx main function kitty protocol', () => { outputFormat: undefined, fakeResponses: undefined, recordResponses: undefined, + rawOutput: undefined, + acceptRawOutputRisk: undefined, }); await act(async () => { diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 8a05ac1178..7b12f864b3 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -173,6 +173,8 @@ describe('runNonInteractive', () => { getModel: vi.fn().mockReturnValue('test-model'), getFolderTrust: vi.fn().mockReturnValue(false), isTrustedFolder: vi.fn().mockReturnValue(false), + getRawOutput: vi.fn().mockReturnValue(false), + getAcceptRawOutputRisk: vi.fn().mockReturnValue(false), } as unknown as Config; mockSettings = { @@ -1745,6 +1747,121 @@ describe('runNonInteractive', () => { // The key assertion: sendMessageStream should have been called ONLY ONCE (initial user input). expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledTimes(1); + + expect(processStderrSpy).toHaveBeenCalledWith( + 'Agent execution stopped: Stop reason from hook\n', + ); + }); + + it('should write JSON output when a tool call returns STOP_EXECUTION error', async () => { + vi.mocked(mockConfig.getOutputFormat).mockReturnValue(OutputFormat.JSON); + vi.mocked(uiTelemetryService.getMetrics).mockReturnValue( + MOCK_SESSION_METRICS, + ); + + const toolCallEvent: ServerGeminiStreamEvent = { + type: GeminiEventType.ToolCallRequest, + value: { + callId: 'stop-call', + name: 'stopTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-stop-json', + }, + }; + + mockCoreExecuteToolCall.mockResolvedValue({ + status: 'error', + request: toolCallEvent.value, + tool: {} as AnyDeclarativeTool, + invocation: {} as AnyToolInvocation, + response: { + callId: 'stop-call', + responseParts: [{ text: 'error occurred' }], + errorType: ToolErrorType.STOP_EXECUTION, + error: new Error('Stop reason'), + resultDisplay: undefined, + }, + }); + + const firstCallEvents: ServerGeminiStreamEvent[] = [ + { type: GeminiEventType.Content, value: 'Partial content' }, + toolCallEvent, + ]; + + mockGeminiClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(firstCallEvents), + ); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Run stop tool', + prompt_id: 'prompt-id-stop-json', + }); + + expect(processStdoutSpy).toHaveBeenCalledWith( + JSON.stringify( + { + session_id: 'test-session-id', + response: 'Partial content', + stats: MOCK_SESSION_METRICS, + }, + null, + 2, + ), + ); + }); + + it('should emit result event when a tool call returns STOP_EXECUTION error in streaming JSON mode', async () => { + vi.mocked(mockConfig.getOutputFormat).mockReturnValue( + OutputFormat.STREAM_JSON, + ); + vi.mocked(uiTelemetryService.getMetrics).mockReturnValue( + MOCK_SESSION_METRICS, + ); + + const toolCallEvent: ServerGeminiStreamEvent = { + type: GeminiEventType.ToolCallRequest, + value: { + callId: 'stop-call', + name: 'stopTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-stop-stream', + }, + }; + + mockCoreExecuteToolCall.mockResolvedValue({ + status: 'error', + request: toolCallEvent.value, + tool: {} as AnyDeclarativeTool, + invocation: {} as AnyToolInvocation, + response: { + callId: 'stop-call', + responseParts: [{ text: 'error occurred' }], + errorType: ToolErrorType.STOP_EXECUTION, + error: new Error('Stop reason'), + resultDisplay: undefined, + }, + }); + + const firstCallEvents: ServerGeminiStreamEvent[] = [toolCallEvent]; + + mockGeminiClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(firstCallEvents), + ); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Run stop tool', + prompt_id: 'prompt-id-stop-stream', + }); + + const output = getWrittenOutput(); + expect(output).toContain('"type":"result"'); + expect(output).toContain('"status":"success"'); }); describe('Agent Execution Events', () => { @@ -1805,4 +1922,142 @@ describe('runNonInteractive', () => { expect(getWrittenOutput()).toBe('Final answer\n'); }); }); + + describe('Output Sanitization', () => { + const ANSI_SEQUENCE = '\u001B[31mRed Text\u001B[0m'; + const OSC_HYPERLINK = + '\u001B]8;;http://example.com\u001B\\Link\u001B]8;;\u001B\\'; + const PLAIN_TEXT_RED = 'Red Text'; + const PLAIN_TEXT_LINK = 'Link'; + + it('should sanitize ANSI output by default', async () => { + const events: ServerGeminiStreamEvent[] = [ + { type: GeminiEventType.Content, value: ANSI_SEQUENCE }, + { type: GeminiEventType.Content, value: ' ' }, + { type: GeminiEventType.Content, value: OSC_HYPERLINK }, + { + type: GeminiEventType.Finished, + value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } }, + }, + ]; + mockGeminiClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(events), + ); + + vi.mocked(mockConfig.getRawOutput).mockReturnValue(false); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Test input', + prompt_id: 'prompt-id-sanitization', + }); + + expect(getWrittenOutput()).toBe(`${PLAIN_TEXT_RED} ${PLAIN_TEXT_LINK}\n`); + }); + + it('should allow ANSI output when rawOutput is true', async () => { + const events: ServerGeminiStreamEvent[] = [ + { type: GeminiEventType.Content, value: ANSI_SEQUENCE }, + { type: GeminiEventType.Content, value: ' ' }, + { type: GeminiEventType.Content, value: OSC_HYPERLINK }, + { + type: GeminiEventType.Finished, + value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } }, + }, + ]; + mockGeminiClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(events), + ); + + vi.mocked(mockConfig.getRawOutput).mockReturnValue(true); + vi.mocked(mockConfig.getAcceptRawOutputRisk).mockReturnValue(true); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Test input', + prompt_id: 'prompt-id-raw', + }); + + expect(getWrittenOutput()).toBe(`${ANSI_SEQUENCE} ${OSC_HYPERLINK}\n`); + }); + + it('should allow ANSI output when only acceptRawOutputRisk is true', async () => { + const events: ServerGeminiStreamEvent[] = [ + { type: GeminiEventType.Content, value: ANSI_SEQUENCE }, + { + type: GeminiEventType.Finished, + value: { reason: undefined, usageMetadata: { totalTokenCount: 5 } }, + }, + ]; + mockGeminiClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(events), + ); + + vi.mocked(mockConfig.getRawOutput).mockReturnValue(false); + vi.mocked(mockConfig.getAcceptRawOutputRisk).mockReturnValue(true); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Test input', + prompt_id: 'prompt-id-accept-only', + }); + + expect(getWrittenOutput()).toBe(`${ANSI_SEQUENCE}\n`); + }); + + it('should warn when rawOutput is true and acceptRisk is false', async () => { + const events: ServerGeminiStreamEvent[] = [ + { + type: GeminiEventType.Finished, + value: { reason: undefined, usageMetadata: { totalTokenCount: 0 } }, + }, + ]; + mockGeminiClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(events), + ); + + vi.mocked(mockConfig.getRawOutput).mockReturnValue(true); + vi.mocked(mockConfig.getAcceptRawOutputRisk).mockReturnValue(false); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Test input', + prompt_id: 'prompt-id-warn', + }); + + expect(processStderrSpy).toHaveBeenCalledWith( + expect.stringContaining('[WARNING] --raw-output is enabled'), + ); + }); + + it('should not warn when rawOutput is true and acceptRisk is true', async () => { + const events: ServerGeminiStreamEvent[] = [ + { + type: GeminiEventType.Finished, + value: { reason: undefined, usageMetadata: { totalTokenCount: 0 } }, + }, + ]; + mockGeminiClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(events), + ); + + vi.mocked(mockConfig.getRawOutput).mockReturnValue(true); + vi.mocked(mockConfig.getAcceptRawOutputRisk).mockReturnValue(true); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Test input', + prompt_id: 'prompt-id-no-warn', + }); + + expect(processStderrSpy).not.toHaveBeenCalledWith( + expect.stringContaining('[WARNING] --raw-output is enabled'), + ); + }); + }); }); diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 8eac2f61d3..50ba2235c4 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -33,6 +33,7 @@ import { import type { Content, Part } from '@google/genai'; import readline from 'node:readline'; +import stripAnsi from 'strip-ansi'; import { convertSessionToHistoryFormats } from './ui/hooks/useSessionBrowser.js'; import { handleSlashCommand } from './nonInteractiveCliCommands.js'; @@ -176,6 +177,16 @@ export async function runNonInteractive({ try { consolePatcher.patch(); + if ( + config.getRawOutput() && + !config.getAcceptRawOutputRisk() && + config.getOutputFormat() === OutputFormat.TEXT + ) { + process.stderr.write( + '[WARNING] --raw-output is enabled. Model output is not sanitized and may contain harmful ANSI sequences (e.g. for phishing or command injection). Use --accept-raw-output-risk to suppress this warning.\n', + ); + } + // Setup stdin cancellation listener setupStdinCancellation(); @@ -285,19 +296,22 @@ export async function runNonInteractive({ } if (event.type === GeminiEventType.Content) { + const isRaw = + config.getRawOutput() || config.getAcceptRawOutputRisk(); + const output = isRaw ? event.value : stripAnsi(event.value); if (streamFormatter) { streamFormatter.emitEvent({ type: JsonStreamEventType.MESSAGE, timestamp: new Date().toISOString(), role: 'assistant', - content: event.value, + content: output, delta: true, }); } else if (config.getOutputFormat() === OutputFormat.JSON) { - responseText += event.value; + responseText += output; } else { if (event.value) { - textOutput.write(event.value); + textOutput.write(output); } } } else if (event.type === GeminiEventType.ToolCallRequest) { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 2871b9fc30..9d21e9c886 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -378,6 +378,8 @@ export interface ConfigParameters { recordResponses?: string; ptyInfo?: string; disableYoloMode?: boolean; + rawOutput?: boolean; + acceptRawOutputRisk?: boolean; modelConfigServiceConfig?: ModelConfigServiceConfig; enableHooks?: boolean; enableHooksUI?: boolean; @@ -520,6 +522,8 @@ export class Config { readonly fakeResponses?: string; readonly recordResponses?: string; private readonly disableYoloMode: boolean; + private readonly rawOutput: boolean; + private readonly acceptRawOutputRisk: boolean; private pendingIncludeDirectories: string[]; private readonly enableHooks: boolean; private readonly enableHooksUI: boolean; @@ -722,6 +726,8 @@ export class Config { }; this.retryFetchErrors = params.retryFetchErrors ?? false; this.disableYoloMode = params.disableYoloMode ?? false; + this.rawOutput = params.rawOutput ?? false; + this.acceptRawOutputRisk = params.acceptRawOutputRisk ?? false; if (params.hooks) { this.hooks = params.hooks; @@ -1395,6 +1401,14 @@ export class Config { return this.disableYoloMode || !this.isTrustedFolder(); } + getRawOutput(): boolean { + return this.rawOutput; + } + + getAcceptRawOutputRisk(): boolean { + return this.acceptRawOutputRisk; + } + getPendingIncludeDirectories(): string[] { return this.pendingIncludeDirectories; }