Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
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
18 changes: 8 additions & 10 deletions packages/cli/src/ui/hooks/shellCommandProcessor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,20 +45,18 @@ function addShellCommandToGeminiHistory(
? resultText.substring(0, MAX_OUTPUT_LENGTH) + '\n... (truncated)'
: resultText;

// Escape backticks to prevent prompt injection breakouts
const safeQuery = rawQuery.replace(/\\/g, '\\\\').replace(/\x60/g, '\\\x60');
const safeModelContent = modelContent
.replace(/\\/g, '\\\\')
.replace(/\x60/g, '\\\x60');

// eslint-disable-next-line @typescript-eslint/no-floating-promises
geminiClient.addHistory({
role: 'user',
parts: [
{
text: `I ran the following shell command:
\`\`\`sh
${rawQuery}
\`\`\`

This produced the following result:
\`\`\`
${modelContent}
\`\`\``,
text: `I ran the following shell command:\n\`\`\`sh\n${safeQuery}\n\`\`\`\n\nThis produced the following result:\n\`\`\`\n${safeModelContent}\n\`\`\``,
},
],
});
Expand Down Expand Up @@ -444,7 +442,7 @@ export const useShellCommandProcessor = (
}

let mainContent: string;
if (isBinary(result.rawOutput)) {
if (isBinaryStream || isBinary(result.rawOutput)) {
mainContent =
'[Command produced binary output, which is not shown.]';
} else {
Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/services/executionLifecycleService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ export type ExecutionMethod =
| 'none';

export interface ExecutionResult {
rawOutput: Buffer;
rawOutput?: Buffer;
output: string;
exitCode: number | null;
signal: number | null;
Expand Down
10 changes: 2 additions & 8 deletions packages/core/src/services/shellExecutionService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -880,15 +880,12 @@ describe('ShellExecutionService', () => {
const binaryChunk1 = Buffer.from([0x89, 0x50, 0x4e, 0x47]);
const binaryChunk2 = Buffer.from([0x0d, 0x0a, 0x1a, 0x0a]);

const { result } = await simulateExecution('cat image.png', (pty) => {
await simulateExecution('cat image.png', (pty) => {
pty.onData.mock.calls[0][0](binaryChunk1);
pty.onData.mock.calls[0][0](binaryChunk2);
pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null });
});

expect(result.rawOutput).toEqual(
Buffer.concat([binaryChunk1, binaryChunk2]),
);
expect(onOutputEventMock).toHaveBeenCalledTimes(4);
expect(onOutputEventMock.mock.calls[0][0]).toEqual({
type: 'binary_detected',
Expand Down Expand Up @@ -1464,15 +1461,12 @@ describe('ShellExecutionService child_process fallback', () => {
const binaryChunk1 = Buffer.from([0x89, 0x50, 0x4e, 0x47]);
const binaryChunk2 = Buffer.from([0x0d, 0x0a, 0x1a, 0x0a]);

const { result } = await simulateExecution('cat image.png', (cp) => {
await simulateExecution('cat image.png', (cp) => {
cp.stdout?.emit('data', binaryChunk1);
cp.stdout?.emit('data', binaryChunk2);
cp.emit('exit', 0, null);
});

expect(result.rawOutput).toEqual(
Buffer.concat([binaryChunk1, binaryChunk2]),
);
expect(onOutputEventMock).toHaveBeenCalledTimes(4);
expect(onOutputEventMock.mock.calls[0][0]).toEqual({
type: 'binary_detected',
Expand Down
46 changes: 26 additions & 20 deletions packages/core/src/services/shellExecutionService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,8 @@ interface ActiveChildProcess {
state: {
output: string;
truncated: boolean;
outputChunks: Buffer[];
sniffChunks: Buffer[];
binaryBytesReceived: number;
};
}

Expand Down Expand Up @@ -493,7 +494,8 @@ export class ShellExecutionService {
const state = {
output: '',
truncated: false,
outputChunks: [] as Buffer[],
sniffChunks: [] as Buffer[],
binaryBytesReceived: 0,
};

if (child.pid) {
Expand Down Expand Up @@ -563,14 +565,19 @@ export class ShellExecutionService {
}
}

state.outputChunks.push(data);
if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) {
state.sniffChunks.push(data);
} else if (!isStreamingRawContent) {
state.binaryBytesReceived += data.length;
}

if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) {
const sniffBuffer = Buffer.concat(state.outputChunks.slice(0, 20));
const sniffBuffer = Buffer.concat(state.sniffChunks.slice(0, 20));
sniffedBytes = sniffBuffer.length;

if (isBinary(sniffBuffer)) {
isStreamingRawContent = false;
state.binaryBytesReceived = sniffBuffer.length;
const event: ShellOutputEvent = { type: 'binary_detected' };
onOutputEvent(event);
if (child.pid) {
Expand Down Expand Up @@ -610,10 +617,7 @@ export class ShellExecutionService {
}
}
} else {
const totalBytes = state.outputChunks.reduce(
(sum, chunk) => sum + chunk.length,
0,
);
const totalBytes = state.binaryBytesReceived;
const event: ShellOutputEvent = {
type: 'binary_progress',
bytesReceived: totalBytes,
Expand All @@ -629,7 +633,7 @@ export class ShellExecutionService {
code: number | null,
signal: NodeJS.Signals | null,
) => {
const { finalBuffer } = cleanup();
cleanup();

let combinedOutput = state.output;
if (state.truncated) {
Expand All @@ -644,7 +648,7 @@ export class ShellExecutionService {
const exitSignal = signal ? os.constants.signals[signal] : null;

const resultPayload: ShellExecutionResult = {
rawOutput: finalBuffer,
rawOutput: Buffer.from(''),
output: finalStrippedOutput,
exitCode,
signal: exitSignal,
Expand Down Expand Up @@ -733,8 +737,7 @@ export class ShellExecutionService {
}
}

const finalBuffer = Buffer.concat(state.outputChunks);
return { finalBuffer };
return;
}

return { pid: child.pid, result };
Expand Down Expand Up @@ -864,7 +867,8 @@ export class ShellExecutionService {
let processingChain = Promise.resolve();
let decoder: TextDecoder | null = null;
let output: string | AnsiOutput | null = null;
const outputChunks: Buffer[] = [];
const sniffChunks: Buffer[] = [];
let binaryBytesReceived = 0;
const error: Error | null = null;
let exited = false;

Expand Down Expand Up @@ -995,14 +999,19 @@ export class ShellExecutionService {
}
}

outputChunks.push(data);
if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) {
sniffChunks.push(data);
} else if (!isStreamingRawContent) {
binaryBytesReceived += data.length;
}

if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) {
const sniffBuffer = Buffer.concat(outputChunks.slice(0, 20));
const sniffBuffer = Buffer.concat(sniffChunks.slice(0, 20));
sniffedBytes = sniffBuffer.length;

if (isBinary(sniffBuffer)) {
isStreamingRawContent = false;
binaryBytesReceived = sniffBuffer.length;
const event: ShellOutputEvent = { type: 'binary_detected' };
onOutputEvent(event);
ExecutionLifecycleService.emitEvent(ptyPid, event);
Expand All @@ -1027,10 +1036,7 @@ export class ShellExecutionService {
resolveChunk();
});
} else {
const totalBytes = outputChunks.reduce(
(sum, chunk) => sum + chunk.length,
0,
);
const totalBytes = binaryBytesReceived;
const event: ShellOutputEvent = {
type: 'binary_progress',
bytesReceived: totalBytes,
Expand Down Expand Up @@ -1076,7 +1082,7 @@ export class ShellExecutionService {
});

ExecutionLifecycleService.completeWithResult(ptyPid, {
rawOutput: Buffer.concat(outputChunks),
rawOutput: Buffer.from(''),
output: getFullBufferText(headlessTerminal),
exitCode,
signal: signal ?? null,
Expand Down
Loading