Adds shell command allowlist (#68)
* Wire through passthrough commands * Add default passthrough commands * Clean up config passing to useGeminiStream
This commit is contained in:
parent
f480ef4bbc
commit
a76d9b4dcf
|
@ -72,6 +72,7 @@ export function loadCliConfig(): Config {
|
|||
argv.model || DEFAULT_GEMINI_MODEL,
|
||||
argv.target_dir || process.cwd(),
|
||||
argv.debug_mode || false,
|
||||
// TODO: load passthroughCommands from .env file
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
@ -32,7 +32,7 @@ export const App = ({ config }: AppProps) => {
|
|||
const [history, setHistory] = useState<HistoryItem[]>([]);
|
||||
const [startupWarnings, setStartupWarnings] = useState<string[]>([]);
|
||||
const { streamingState, submitQuery, initError, debugMessage } =
|
||||
useGeminiStream(setHistory, config.getApiKey(), config.getModel());
|
||||
useGeminiStream(setHistory, config);
|
||||
const { elapsedTime, currentLoadingPhrase } =
|
||||
useLoadingIndicator(streamingState);
|
||||
|
||||
|
|
|
@ -14,6 +14,7 @@ import {
|
|||
getErrorMessage,
|
||||
isNodeError,
|
||||
ToolResult,
|
||||
Config,
|
||||
} from '@gemini-code/server';
|
||||
import type { Chat, PartListUnion, FunctionDeclaration } from '@google/genai';
|
||||
// Import CLI types
|
||||
|
@ -27,8 +28,6 @@ import { StreamingState } from '../../core/gemini-stream.js';
|
|||
// Import CLI tool registry
|
||||
import { toolRegistry } from '../../tools/tool-registry.js';
|
||||
|
||||
const _allowlistedCommands = ['ls']; // Prefix with underscore since it's unused
|
||||
|
||||
const addHistoryItem = (
|
||||
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
|
||||
itemData: Omit<HistoryItem, 'id'>,
|
||||
|
@ -43,8 +42,7 @@ const addHistoryItem = (
|
|||
// Hook now accepts apiKey and model
|
||||
export const useGeminiStream = (
|
||||
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
|
||||
apiKey: string,
|
||||
model: string,
|
||||
config: Config,
|
||||
) => {
|
||||
const [streamingState, setStreamingState] = useState<StreamingState>(
|
||||
StreamingState.Idle,
|
||||
|
@ -62,15 +60,17 @@ export const useGeminiStream = (
|
|||
setInitError(null);
|
||||
if (!geminiClientRef.current) {
|
||||
try {
|
||||
geminiClientRef.current = new GeminiClient(apiKey, model);
|
||||
geminiClientRef.current = new GeminiClient(
|
||||
config.getApiKey(),
|
||||
config.getModel(),
|
||||
);
|
||||
} catch (error: unknown) {
|
||||
setInitError(
|
||||
`Failed to initialize client: ${getErrorMessage(error) || 'Unknown error'}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
// Dependency array includes apiKey and model now
|
||||
}, [apiKey, model]);
|
||||
}, [config.getApiKey(), config.getModel()]);
|
||||
|
||||
// Input Handling Effect (remains the same)
|
||||
useInput((input, key) => {
|
||||
|
@ -107,6 +107,39 @@ export const useGeminiStream = (
|
|||
|
||||
if (typeof query === 'string') {
|
||||
setDebugMessage(`User query: ${query}`);
|
||||
const maybeCommand = query.split(/\s+/)[0];
|
||||
if (config.getPassthroughCommands().includes(maybeCommand)) {
|
||||
// Execute and capture output
|
||||
setDebugMessage(`Executing shell command directly: ${query}`);
|
||||
_exec(query, (error, stdout, stderr) => {
|
||||
const timestamp = getNextMessageId(Date.now());
|
||||
if (error) {
|
||||
addHistoryItem(
|
||||
setHistory,
|
||||
{ type: 'error', text: error.message },
|
||||
timestamp,
|
||||
);
|
||||
} else if (stderr) {
|
||||
addHistoryItem(
|
||||
setHistory,
|
||||
{ type: 'error', text: stderr },
|
||||
timestamp,
|
||||
);
|
||||
} else {
|
||||
// Add stdout as an info message
|
||||
addHistoryItem(
|
||||
setHistory,
|
||||
{ type: 'info', text: stdout || '' },
|
||||
timestamp,
|
||||
);
|
||||
}
|
||||
// Set state back to Idle *after* command finishes and output is added
|
||||
setStreamingState(StreamingState.Idle);
|
||||
});
|
||||
// Set state to Responding while the command runs
|
||||
setStreamingState(StreamingState.Responding);
|
||||
return; // Prevent Gemini call
|
||||
}
|
||||
}
|
||||
|
||||
const userMessageTimestamp = Date.now();
|
||||
|
@ -391,7 +424,8 @@ export const useGeminiStream = (
|
|||
}
|
||||
} finally {
|
||||
abortControllerRef.current = null;
|
||||
// Only set to Idle if not waiting for confirmation
|
||||
// Only set to Idle if not waiting for confirmation.
|
||||
// Passthrough commands handle their own Idle transition.
|
||||
if (streamingState !== StreamingState.WaitingForConfirmation) {
|
||||
setStreamingState(StreamingState.Idle);
|
||||
}
|
||||
|
@ -401,8 +435,8 @@ export const useGeminiStream = (
|
|||
[
|
||||
streamingState,
|
||||
setHistory,
|
||||
apiKey,
|
||||
model,
|
||||
config.getApiKey(),
|
||||
config.getModel(),
|
||||
getNextMessageId,
|
||||
updateGeminiMessage,
|
||||
],
|
||||
|
|
|
@ -9,22 +9,28 @@ import * as fs from 'node:fs';
|
|||
import * as path from 'node:path';
|
||||
import process from 'node:process';
|
||||
|
||||
const DEFAULT_PASSTHROUGH_COMMANDS = ['ls', 'git', 'npm'];
|
||||
|
||||
export class Config {
|
||||
private apiKey: string;
|
||||
private model: string;
|
||||
private targetDir: string;
|
||||
private debugMode: boolean;
|
||||
private passthroughCommands: string[];
|
||||
|
||||
constructor(
|
||||
apiKey: string,
|
||||
model: string,
|
||||
targetDir: string,
|
||||
debugMode: boolean,
|
||||
passthroughCommands?: string[],
|
||||
) {
|
||||
this.apiKey = apiKey;
|
||||
this.model = model;
|
||||
this.targetDir = targetDir;
|
||||
this.debugMode = debugMode;
|
||||
this.passthroughCommands =
|
||||
passthroughCommands || DEFAULT_PASSTHROUGH_COMMANDS;
|
||||
}
|
||||
|
||||
getApiKey(): string {
|
||||
|
@ -42,6 +48,10 @@ export class Config {
|
|||
getDebugMode(): boolean {
|
||||
return this.debugMode;
|
||||
}
|
||||
|
||||
getPassthroughCommands(): string[] {
|
||||
return this.passthroughCommands;
|
||||
}
|
||||
}
|
||||
|
||||
function findEnvFile(startDir: string): string | null {
|
||||
|
@ -72,6 +82,13 @@ export function createServerConfig(
|
|||
model: string,
|
||||
targetDir: string,
|
||||
debugMode: boolean,
|
||||
passthroughCommands?: string[],
|
||||
): Config {
|
||||
return new Config(apiKey, model, path.resolve(targetDir), debugMode);
|
||||
return new Config(
|
||||
apiKey,
|
||||
model,
|
||||
path.resolve(targetDir),
|
||||
debugMode,
|
||||
passthroughCommands,
|
||||
);
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue