z-bar-qt/Plugins/ZShell/Llm/shelltool.cpp

#include "shelltool.hpp"

#include <QDir>
#include <QJsonArray>
#include <QJsonObject>
#include <QProcess>
#include <QTimer>

#include <unistd.h>

namespace ZShell::llm {

ShellTool::ShellTool(ToolRegistry* registry, QObject* parent)
	: LlmTool(registry, parent) {}

ShellTool::~ShellTool() {
	for (auto* job : m_jobs) {
		if (job->timer) job->timer->stop();
		if (job->process) job->process->kill();
		delete job;
	}
	m_jobs.clear();
}

QString ShellTool::name() const {
	return QStringLiteral("runcommand");
}

QString ShellTool::description() const {
	return QStringLiteral(
		"Run a shell command inside a bubblewrap sandbox that contains "
		"the files the user attached to the conversation. The rest of "
		"the filesystem is mounted read-only, there is no network "
		"access, and only the sandbox directory is writable. The "
		"command is executed with /bin/sh and the sandbox directory as "
		"its working directory, so relative paths refer to the sandbox. "
		"Use it to explore, transform, or extract attached files (for "
		"example: ls, find, grep, unzip, tar, head, wc). The combined "
		"output is limited; the exit code is always reported. The call "
		"only completes when every child process has exited, so do not "
		"start long-lived background processes.");
}

QJsonObject ShellTool::parameters() const {
	QJsonObject command;
	command[QStringLiteral("type")] = QStringLiteral("string");
	command[QStringLiteral("description")] =
		QStringLiteral("The shell command to execute");

	QJsonObject timeout;
	timeout[QStringLiteral("type")] = QStringLiteral("integer");
	timeout[QStringLiteral("minimum")] = 1;
	timeout[QStringLiteral("maximum")] = MaxTimeoutSeconds;
	timeout[QStringLiteral("description")] =
		QStringLiteral("Timeout in seconds. Defaults to %1.")
			.arg(DefaultTimeoutSeconds);

	QJsonObject properties;
	properties[QStringLiteral("command")] = command;
	properties[QStringLiteral("timeout")] = timeout;

	QJsonObject schema;
	schema[QStringLiteral("type")] = QStringLiteral("object");
	schema[QStringLiteral("properties")] = properties;
	QJsonArray required;
	required.append(QStringLiteral("command"));
	schema[QStringLiteral("required")] = required;
	return schema;
}

QStringList ShellTool::sandboxArguments(const QString& workspace) {
	QStringList args;
	args << QLatin1String("--unshare-user") << QLatin1String("--uid")
		 << QString::number(::getuid()) << QLatin1String("--gid")
		 << QString::number(::getgid()) << QLatin1String("--disable-userns")
		 << QLatin1String("--unshare-pid") << QLatin1String("--unshare-net")
		 << QLatin1String("--unshare-ipc") << QLatin1String("--unshare-uts")
		 << QLatin1String("--hostname") << QLatin1String("zshell-sandbox");
	args << QLatin1String("--ro-bind") << QLatin1String("/")
		 << QLatin1String("/") << QLatin1String("--tmpfs")
		 << QLatin1String("/tmp") << QLatin1String("--tmpfs")
		 << QLatin1String("/run") << QLatin1String("--bind") << workspace
		 << workspace << QLatin1String("--dev") << QLatin1String("/dev")
		 << QLatin1String("--proc") << QLatin1String("/proc");
	args << QLatin1String("--chdir") << workspace;
	args << QLatin1String("--clearenv") << QLatin1String("--setenv")
		 << QLatin1String("PATH")
		 << QLatin1String(
				"/usr/local/sbin:/usr/local/bin:/usr/sbin:"
				"/usr/bin:/sbin:/bin")
		 << QLatin1String("--setenv") << QLatin1String("HOME") << workspace
		 << QLatin1String("--setenv") << QLatin1String("TMPDIR")
		 << QLatin1String("/tmp") << QLatin1String("--setenv")
		 << QLatin1String("SHELL") << QLatin1String("/bin/sh");
	const QByteArray user = qgetenv("USER");
	if (!user.isEmpty())
		args << QLatin1String("--setenv") << QLatin1String("USER")
			 << QString::fromLocal8Bit(user) << QLatin1String("--setenv")
			 << QLatin1String("LOGNAME") << QString::fromLocal8Bit(user);
	const QByteArray lang = qgetenv("LANG");
	if (!lang.isEmpty())
		args << QLatin1String("--setenv") << QLatin1String("LANG")
			 << QString::fromLocal8Bit(lang);
	args << QLatin1String("--new-session")
		 << QLatin1String("--die-with-parent");
	return args;
}

bool ShellTool::available(const ToolContext& context) const {
	return registry()->sandboxEnabled() && !context.workspace.isEmpty();
}

void ShellTool::execute(
	const ToolContext& context, const QJsonObject& args, ToolDone done) {
	const QString command = args[QStringLiteral("command")].toString();
	if (command.trimmed().isEmpty()) {
		done({false, QStringLiteral("Missing required argument 'command'")});
		return;
	}
	if (!registry()->sandboxEnabled()) {
		done(
			{false,
			 QStringLiteral(
				 "The file sandbox is disabled, so no commands can run")});
		return;
	}
	const QString workspace = context.workspace;
	if (workspace.isEmpty()) {
		done(
			{false,
			 QStringLiteral("No sandbox available: no chat session is active")});
		return;
	}
	QDir().mkpath(workspace);

	const int timeoutSeconds = qBound(
		1,
		args[QStringLiteral("timeout")].toInt(DefaultTimeoutSeconds),
		MaxTimeoutSeconds);
	auto* job = new Job;
	job->timeoutSeconds = timeoutSeconds;
	job->outputBudgetChars = context.outputBudgetChars;
	job->done = std::move(done);
	m_jobs.append(job);

	QProcess* process = new QProcess(this);
	process->setWorkingDirectory(workspace);
	process->setProcessChannelMode(QProcess::MergedChannels);
	job->process = process;

	QTimer* timer = new QTimer(this);
	timer->setSingleShot(true);
	connect(timer, &QTimer::timeout, this, [this, job, process]() {
		job->timedOut = true;
		process->kill();
	});
	timer->start(timeoutSeconds * 1000);
	job->timer = timer;

	connect(
		process,
		QOverload<int, QProcess::ExitStatus>::of(&QProcess::finished),
		this,
		[this, job](int, QProcess::ExitStatus) { finishJob(job); });
	connect(
		process,
		&QProcess::errorOccurred,
		this,
		[this, job, process](QProcess::ProcessError error) {
			if (job->finished) return;
			if (error == QProcess::FailedToStart ||
				process->state() == QProcess::NotRunning) {
				job->startError = process->errorString();
				finishJob(job);
			} else {
				process->kill();
			}
		});

	QStringList commandArgs = sandboxArguments(workspace);
	commandArgs << QStringLiteral("sh") << QStringLiteral("-c") << command;
	process->start(QStringLiteral("bwrap"), commandArgs);
}

void ShellTool::cancel() {
	for (auto* job : m_jobs) {
		if (job->finished || !job->process) continue;
		job->cancelled = true;
		job->process->kill();
	}
}

void ShellTool::finishJob(Job* job) {
	if (!job || job->finished) return;
	job->finished = true;
	if (job->timer) job->timer->stop();
	m_jobs.removeAll(job);

	QProcess* process = job->process;
	const bool ran = process && process->state() != QProcess::NotRunning;
	const int exitCode = ran ? process->exitCode() : -1;
	QString output;
	if (process) output = QString::fromUtf8(process->readAllStandardOutput());

	QString note;
	if (!job->startError.isEmpty())
		note =
			QStringLiteral("Command failed to start: %1\n").arg(job->startError);
	else if (job->cancelled)
		note = QStringLiteral("Command was cancelled.\n");
	else if (job->timedOut)
		note = QStringLiteral(
				   "Command timed out after %1 seconds and was killed.\n")
				   .arg(job->timeoutSeconds);

	int budget = job->outputBudgetChars > 0 ? job->outputBudgetChars
											: registry()->outputBudgetChars(1);
	if (output.size() > budget)
		output = output.left(budget) + QStringLiteral(
										   "\n[output truncated: %1 characters "
										   "total]")
										   .arg(output.size());

	const QString result =
		note + QStringLiteral("exit code: %1\n").arg(exitCode) +
		(output.isEmpty() ? QStringLiteral("(no output)") : output);
	const bool ok = !job->cancelled && !job->timedOut && exitCode == 0;

	ToolDone done = std::move(job->done);
	if (process) process->deleteLater();
	delete job;
	if (done) done({ok, result});
}

} // namespace ZShell::llm