diff --git a/configs/preferences.default.yaml b/configs/preferences.default.yaml index 725c6c7..33094fd 100644 --- a/configs/preferences.default.yaml +++ b/configs/preferences.default.yaml @@ -5,6 +5,7 @@ summarizerModel: gpt-4.1-mini relevanceModel: gpt-4.1-mini webTranslatorModel: gpt-4.1-mini shellCommand: zsh +containerImage: alpine agentConfig: maxIterationLimit: 50 maxRuntimeMsLimit: 300000 # 5m diff --git a/src/chat/context.ts b/src/chat/context.ts index 1126069..df70a45 100644 --- a/src/chat/context.ts +++ b/src/chat/context.ts @@ -5,6 +5,7 @@ import { ChatProviders } from '../providers/chat_providers' import { EmbeddingsProviders } from '../providers/embeddings_providers' import { Preferences } from '../providers/preferences' import { Rule } from '../rules/types' +import { Container } from '../util/docker/container' import { InterruptHandler } from '../util/interrupts/interrupts' import { Prompter } from '../util/prompter/prompter' import { UsageTracker } from '../util/usage/tracker' @@ -21,6 +22,7 @@ export type ChatContext = { contextStateManager: ContextStateManager yolo: boolean tools: string[] + container?: Container } export async function swapProvider(context: ChatContext, modelName: string): Promise { diff --git a/src/cli.ts b/src/cli.ts index 63cf818..3460fe5 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -18,6 +18,7 @@ import { Rule } from './rules/types' import { buildSystemPrompt } from './system' import { seedAllowedCommands } from './tools/shell/shell_execute' import { removeDisabledTools } from './tools/tools' +import { startContainer } from './util/docker/container' import { createInterruptHandler, InterruptHandlerOptions } from './util/interrupts/interrupts' import { createPrompter } from './util/prompter/prompter' import { createLimiter } from './util/ratelimits/limiter' @@ -61,6 +62,9 @@ async function main() { const disableToolsFlags = '--disable-tools ' const disableToolsDescription = 'Comma-separated list of tool names to disable.' + const sandboxedFlags = '--sandboxed [container]' + const sandboxedDescription = 'Run shell commands in a Docker container.' + program .option(historyFlags, historyDescription) .option(portFlags, portDescription) @@ -68,6 +72,7 @@ async function main() { .option(yoloFlags, yoloDescription) .option(oneShotFlags, oneShotDescription) .option(disableToolsFlags, disableToolsDescription) + .option(sandboxedFlags, sandboxedDescription) .action(options => { if (options.cwd) { process.chdir(options.cwd) @@ -84,6 +89,7 @@ async function main() { options.oneShot, options.yolo, options.disableTools, + options.sandboxed, ) }) @@ -101,6 +107,7 @@ async function chat( oneShot?: string, yolo: boolean = false, disabledTools?: string, + sandboxed: string | boolean = false, ) { if (!process.stdin.setRawMode) { throw new Error('chat command is not supported in this environment.') @@ -109,6 +116,7 @@ async function chat( let context: ChatContext const allowedTools = removeDisabledTools(disabledTools?.split(',').map(name => name.trim()) ?? [], 'main') const allowedToolNames = allowedTools.map(({ name }) => name) + const containerImage = sandboxed === true ? preferences.containerImage : (sandboxed ?? '') readline.emitKeypressEvents(process.stdin) process.stdin.setRawMode(true) @@ -136,14 +144,19 @@ async function chat( await registerTools(client) await seedAllowedCommands() + const interruptHandler = createInterruptHandler(rl) + const interruptInputOptions = rootInterruptHandlerOptions(rl) + + const execContext = { interruptHandler, preferences, contextStateManager } + const container = containerImage ? await startContainer(execContext, containerImage) : undefined + try { - const interruptHandler = createInterruptHandler(rl) - const interruptInputOptions = rootInterruptHandlerOptions(rl) const prompter = createPrompter( rl, interruptHandler, attentionGetter(preferences.attentionCommand ?? defaultAttentionCommand), ) + const provider = await providers.createProvider({ contextState: contextStateManager, modelName: preferences.defaultModel, @@ -164,6 +177,7 @@ async function chat( contextStateManager, yolo, tools: allowedToolNames, + container, } await registerContextListeners(context, client) @@ -180,6 +194,7 @@ async function chat( interruptInputOptions, ) } finally { + await container?.stop() process.stdin.unpipe(filter) filter.destroy() rl.close() diff --git a/src/providers/preferences.ts b/src/providers/preferences.ts index 9b32fc0..0b2e5ea 100644 --- a/src/providers/preferences.ts +++ b/src/providers/preferences.ts @@ -56,6 +56,7 @@ const PreferencesSchema = z.object({ z.array(EmbeddingModelSchema), ), shellCommand: z.enum(['bash', 'zsh', 'fish']).optional(), + containerImage: z.string(), attentionCommand: z.string().optional(), agentConfig: z .object({ diff --git a/src/system.ts b/src/system.ts index 6afb737..ed9b352 100644 --- a/src/system.ts +++ b/src/system.ts @@ -152,16 +152,31 @@ Remember: Begin your assistance by analyzing the user's query and providing an appropriate response. +{{environment}} + +{{custom instructions}} +` + +const hostEnvironmentTemplate = ` The current directory is {{cwd}}. The user's preferred shell is {{shellCommand}}. +` -{{custom instructions}} +const containerEnvironmentTemplate = ` +The current directory is /workspace. +You are in a Docker container running {{image}}. +Commands you run will be invoked via sh -c '...'. ` -export async function buildSystemPrompt(preferences: Preferences): Promise { +export async function buildSystemPrompt(preferences: Preferences, containerImage?: string): Promise { + const environment = containerImage + ? containerEnvironmentTemplate.replace('{{image}}', containerImage) + : hostEnvironmentTemplate + .replace('{{cwd}}', process.cwd()) + .replace('{{shellCommand}}', preferences.shellCommand ?? 'zsh') + return systemPromptTemplate - .replace('{{cwd}}', process.cwd()) - .replace('{{shellCommand}}', preferences.shellCommand ?? 'zsh') + .replace('{{environment}}', environment) .replace('{{custom instructions}}', await buildProjectInstructions()) } diff --git a/src/util/docker/container.ts b/src/util/docker/container.ts new file mode 100644 index 0000000..b411200 --- /dev/null +++ b/src/util/docker/container.ts @@ -0,0 +1,41 @@ +import chalk from 'chalk' +import { ExecContext, execOnHost } from '../shell/exec' + +export type Container = Awaited> + +export async function startContainer(context: ExecContext, containerImage: string) { + console.log(`${chalk.dim('ℹ')} Starting Docker container from ${containerImage}...`) + + let containerId = '' + try { + const output = await execOnHost( + context, + `docker run -d -v "${process.cwd()}:/workspace" -w /workspace ${containerImage} tail -f /dev/null`, + ) + + for (const line of output) { + containerId = line.content.trim() + } + + console.log(`${chalk.green('✓')} Docker container ${containerId} started`) + } catch (error: any) { + throw new Error(`Failed to start Docker container: ${error.message}`) + } + + const stop = async () => { + console.log(`${chalk.dim('ℹ')} Stopping Docker container...`) + + try { + await execOnHost(context, `docker stop ${containerId}`) + await execOnHost(context, `docker rm ${containerId}`) + console.log(`${chalk.green('✓')} Docker container stopped and removed`) + } catch (error: any) { + console.warn(`${chalk.yellow('⚠')} Failed to clean up container: ${error.message}`) + } + } + + return { + containerId: () => containerId, + stop, + } +} diff --git a/src/util/shell/exec.ts b/src/util/shell/exec.ts index 692e9c4..b0fe6ea 100644 --- a/src/util/shell/exec.ts +++ b/src/util/shell/exec.ts @@ -15,11 +15,14 @@ export type OutputLine = { content: string } +export type ExecContext = Pick + export async function executeCommand( - context: ChatContext, + context: ExecContext, command: string, + signal?: AbortSignal, // TODO - check callers ): Promise<{ result?: ShellResult; error?: Error; canceled?: boolean }> { - const response = await withProgress(update => runCommand(context, command, update), { + const response = await withProgress(update => runCommand(context, command, update, signal), { progress: prefixFormatter('Executing command...', formatOutput), success: prefixFormatter('Command succeeded.', formatOutput), failure: prefixFormatter('Command failed.', formatOutput), @@ -43,7 +46,26 @@ export async function executeCommand( } } -async function runCommand(context: ChatContext, command: string, update: Updater): Promise { +async function runCommand( + context: ExecContext, + command: string, + update: Updater = () => {}, + signal?: AbortSignal, +): Promise { + const containerId = context.container?.containerId + if (containerId) { + command = `docker exec ${containerId} sh -c ${command}` + } + + return execOnHost(context, command, update, signal) +} + +export async function execOnHost( + context: ExecContext, + command: string, + update: Updater = () => {}, + _signal?: AbortSignal, // TODO - use +): Promise { return context.interruptHandler.withInterruptHandler(signal => { return new Promise((resolve, reject) => { const output: OutputLine[] = []