Skip to content
Open
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
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2025 Google LLC
* Copyright 2026 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand All @@ -19,33 +19,81 @@
import static java.util.Objects.requireNonNullElse;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
import com.github.dockerjava.api.model.Container;
import com.github.dockerjava.api.model.Capability;
import com.github.dockerjava.api.model.HostConfig;
import com.github.dockerjava.core.DefaultDockerClientConfig;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.command.ExecStartResultCallback;
import com.google.adk.agents.InvocationContext;
import com.google.adk.codeexecutors.CodeExecutionUtils.CodeExecutionInput;
import com.google.adk.codeexecutors.CodeExecutionUtils.CodeExecutionResult;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Paths;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** A code executor that uses a custom container to execute code. */
public class ContainerCodeExecutor extends BaseCodeExecutor {
/**
* A code executor that runs code in a Docker container.
*
* <p>A fresh container is created for every {@link #executeCode} call and force-removed afterwards,
* so code from one execution cannot observe or affect another's environment. Code is run via {@code
* docker exec} (as in ADK Python), so the image only needs {@code python3} on its PATH; any image
* {@code ENTRYPOINT} is bypassed.
*
* <p><b>Sandboxing is opt-in.</b> By default the execution container is unrestricted (network
* enabled, writable filesystem, no resource or time limits), matching the previous behavior so
* existing callers are not broken; a warning is logged when it is used this way. Call {@link
* #setStrictSandbox(boolean) setStrictSandbox(true)} to harden each container: no network (unless
* re-enabled via {@link #setNetworkEnabled(boolean)}), all Linux capabilities dropped, no privilege
* escalation, a read-only root filesystem with a small writable {@code /tmp} tmpfs, memory/PID
* limits, and a wall-clock execution timeout. Strict sandboxing becomes the default in ADK 2.0.
*
* <p>The execution timeout and memory limit used by the strict sandbox are configurable via {@link
* #setExecutionTimeoutSeconds(long)} and {@link #setMemoryLimitBytes(long)}.
*
* <p>This executor holds a {@link DockerClient}; call {@link #close()} (or rely on the registered
* JVM shutdown hook) to release its connections and threads. As with ADK Python, an abrupt JVM
* termination (e.g. SIGKILL) during an execution may leave a container behind.
*/
public class ContainerCodeExecutor extends BaseCodeExecutor implements AutoCloseable {
private static final Logger logger = LoggerFactory.getLogger(ContainerCodeExecutor.class);
private static final String DEFAULT_IMAGE_TAG = "adk-code-executor:latest";

/** Default memory limit for each execution container (512 MiB). */
private static final long DEFAULT_MEMORY_LIMIT_BYTES = 512L * 1024 * 1024;

/** Maximum number of processes/threads allowed inside an execution container. */
private static final long PIDS_LIMIT = 128L;

/** Default max wall-clock time a single execution may run before its container is killed. */
private static final long DEFAULT_EXECUTION_TIMEOUT_SECONDS = 60L;

private final String baseUrl;
private final String image;
private final String dockerPath;
private final DockerClient dockerClient;
private Container container;
// Registered by the image/dockerPath constructor as a backstop; removed in close() so a closed
// executor is not retained by the JVM's shutdown-hook list.
private final Thread shutdownHook = new Thread(this::close);
private boolean networkEnabled = false;
private long executionTimeoutSeconds = DEFAULT_EXECUTION_TIMEOUT_SECONDS;
private long memoryLimitBytes = DEFAULT_MEMORY_LIMIT_BYTES;

// Off by default so this executor does not change behavior for existing callers; a warning is
// logged while it is disabled, and it becomes the default in ADK 2.0.
private boolean strictSandbox = false;
private final AtomicBoolean strictSandboxWarningLogged = new AtomicBoolean(false);

/**
* Creates a ContainerCodeExecutor from an image.
Expand Down Expand Up @@ -99,17 +147,65 @@ public ContainerCodeExecutor(String baseUrl, String image, String dockerPath) {
this.baseUrl = baseUrl;
this.image = requireNonNullElse(image, DEFAULT_IMAGE_TAG);
this.dockerPath = dockerPath == null ? null : Paths.get(dockerPath).toAbsolutePath().toString();
this.dockerClient = buildDockerClient(baseUrl);
prepareImage();
// Backstop so the client is released even if callers forget to close() this executor.
Runtime.getRuntime().addShutdownHook(shutdownHook);
}

if (baseUrl != null) {
var config =
DefaultDockerClientConfig.createDefaultConfigBuilder().withDockerHost(baseUrl).build();
this.dockerClient = DockerClientBuilder.getInstance(config).build();
} else {
this.dockerClient = DockerClientBuilder.getInstance().build();
}
/** Test-only constructor that injects a Docker client and skips image preparation. */
@VisibleForTesting
ContainerCodeExecutor(DockerClient dockerClient, String image) {
this.baseUrl = null;
this.image = requireNonNullElse(image, DEFAULT_IMAGE_TAG);
this.dockerPath = null;
this.dockerClient = dockerClient;
}

/**
* Enables or disables container networking when the strict sandbox is on. In strict mode
* networking is disabled by default so executed code cannot reach the network (including the
* cloud metadata endpoint); pass {@code true} to allow it. Has no effect unless {@link
* #setStrictSandbox(boolean)} is enabled — without the sandbox the container always has network
* access.
*/
public ContainerCodeExecutor setNetworkEnabled(boolean networkEnabled) {
this.networkEnabled = networkEnabled;
return this;
}

/**
* Sets the maximum wall-clock time (in seconds) a single execution may run, in the strict
* sandbox, before its container is force-removed (killed). Defaults to 60 seconds. Has no effect
* unless {@link #setStrictSandbox(boolean)} is enabled.
*/
public ContainerCodeExecutor setExecutionTimeoutSeconds(long executionTimeoutSeconds) {
this.executionTimeoutSeconds = executionTimeoutSeconds;
return this;
}

/**
* Sets the per-execution container memory limit, in bytes, used by the strict sandbox. Defaults
* to 512 MiB. Has no effect unless {@link #setStrictSandbox(boolean)} is enabled.
*/
public ContainerCodeExecutor setMemoryLimitBytes(long memoryLimitBytes) {
this.memoryLimitBytes = memoryLimitBytes;
return this;
}

initContainer();
Runtime.getRuntime().addShutdownHook(new Thread(this::cleanupContainer));
/**
* Enables the strict sandbox. When enabled, each execution runs in a hardened container: no
* network (unless re-enabled via {@link #setNetworkEnabled(boolean)}), all Linux capabilities
* dropped, no privilege escalation, a read-only root filesystem (writable {@code /tmp} only),
* memory/PID limits, and a wall-clock timeout.
*
* <p>Disabled by default so enabling the sandbox is not a breaking change for existing callers.
* While it is disabled a warning is logged, because running untrusted, model-generated code
* without the sandbox is dangerous. Strict sandboxing becomes the default in ADK 2.0.
*/
public ContainerCodeExecutor setStrictSandbox(boolean strictSandbox) {
this.strictSandbox = strictSandbox;
return this;
}

@Override
Expand All @@ -125,70 +221,140 @@ public boolean optimizeDataFile() {
@Override
public CodeExecutionResult executeCode(
InvocationContext invocationContext, CodeExecutionInput codeExecutionInput) {
warnIfStrictSandboxDisabled();

ByteArrayOutputStream stdout = new ByteArrayOutputStream();
ByteArrayOutputStream stderr = new ByteArrayOutputStream();

ExecCreateCmdResponse execCreateCmdResponse =
dockerClient
.execCreateCmd(container.getId())
.withAttachStdout(true)
.withAttachStderr(true)
.withCmd("python3", "-c", codeExecutionInput.code())
.exec();
// A fresh container per execution isolates each run from every other session's execution
// environment. Code is run via `docker exec` (as in ADK Python), which needs only `python3` on
// the image and bypasses any ENTRYPOINT. The hardened HostConfig is only applied in the strict
// sandbox; otherwise the container is left unrestricted to preserve existing behavior.
var createContainerCmd =
dockerClient.createContainerCmd(image).withTty(true).withAttachStdin(true);
if (strictSandbox) {
createContainerCmd.withHostConfig(sandboxHostConfig());
}
CreateContainerResponse createContainerResponse = createContainerCmd.exec();
String containerId = createContainerResponse.getId();
try {
dockerClient
.execStartCmd(execCreateCmdResponse.getId())
.exec(new ExecStartResultCallback(stdout, stderr))
.awaitCompletion();
dockerClient.startContainerCmd(containerId).exec();

ExecCreateCmdResponse execCreateCmdResponse =
dockerClient
.execCreateCmd(containerId)
.withAttachStdout(true)
.withAttachStderr(true)
.withCmd("python3", "-c", codeExecutionInput.code())
.exec();

boolean completed;
try (ExecStartResultCallback callback = new ExecStartResultCallback(stdout, stderr)) {
dockerClient.execStartCmd(execCreateCmdResponse.getId()).exec(callback);
if (strictSandbox) {
completed = callback.awaitCompletion(executionTimeoutSeconds, TimeUnit.SECONDS);
} else {
// No execution timeout unless the strict sandbox is enabled, matching prior behavior.
callback.awaitCompletion();
completed = true;
}
}

if (!completed) {
// Force-removing the container in the finally block kills the still-running execution.
return CodeExecutionResult.builder()
.stderr(
String.format(
"Code execution timed out after %d seconds.", executionTimeoutSeconds))
.build();
}
return CodeExecutionResult.builder()
.stdout(stdout.toString(StandardCharsets.UTF_8))
.stderr(stderr.toString(StandardCharsets.UTF_8))
.build();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Code execution was interrupted.", e);
} catch (IOException e) {
throw new UncheckedIOException(e);
} finally {
removeContainerQuietly(containerId);
}

return CodeExecutionResult.builder()
.stdout(stdout.toString(StandardCharsets.UTF_8))
.stderr(stderr.toString(StandardCharsets.UTF_8))
.build();
}

private void buildDockerImage() {
if (dockerPath == null) {
throw new IllegalStateException("Docker path is not set.");
/** Builds the hardened {@link HostConfig} applied to each execution container in strict mode. */
@VisibleForTesting
HostConfig sandboxHostConfig() {
HostConfig hostConfig =
HostConfig.newHostConfig()
.withCapDrop(Capability.ALL)
.withReadonlyRootfs(true)
.withSecurityOpts(ImmutableList.of("no-new-privileges"))
.withMemory(memoryLimitBytes)
.withPidsLimit(PIDS_LIMIT)
// A read-only rootfs still needs a small writable scratch space at /tmp.
.withTmpFs(ImmutableMap.of("/tmp", "rw,size=64m"));
if (!networkEnabled) {
hostConfig.withNetworkMode("none");
}
File dockerfile = new File(dockerPath);
if (!dockerfile.exists()) {
throw new UncheckedIOException(new IOException("Invalid Docker path: " + dockerPath));
return hostConfig;
}

/**
* Logs a warning, at most once per executor, if the strict sandbox is disabled. Returns whether
* the warning was logged.
*/
@VisibleForTesting
boolean warnIfStrictSandboxDisabled() {
if (!strictSandbox && strictSandboxWarningLogged.compareAndSet(false, true)) {
logger.warn(
"ContainerCodeExecutor is running with the strict sandbox disabled (the current default):"
+ " the execution container has network access (including the cloud metadata"
+ " endpoint), a writable filesystem, and no memory/PID/time limits, so untrusted,"
+ " model-generated code can steal credentials, exhaust host resources, or hang the"
+ " calling thread. Call setStrictSandbox(true) to run each execution in a"
+ " locked-down container. This becomes the default in ADK 2.0.");
return true;
}
return false;
}

logger.info("Building Docker image...");
private void removeContainerQuietly(String containerId) {
try {
dockerClient.buildImageCmd(dockerfile).withTag(image).start().awaitCompletion();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Docker image build was interrupted.", e);
dockerClient.removeContainerCmd(containerId).withForce(true).exec();
} catch (RuntimeException e) {
logger.warn("Failed to remove container {}", containerId, e);
}
logger.info("Docker image: {} built.", image);
}

private void verifyPythonInstallation() {
ExecCreateCmdResponse execCreateCmdResponse =
dockerClient.execCreateCmd(container.getId()).withCmd("which", "python3").exec();
ByteArrayOutputStream stdout = new ByteArrayOutputStream();
ByteArrayOutputStream stderr = new ByteArrayOutputStream();
try (ExecStartResultCallback callback = new ExecStartResultCallback(stdout, stderr)) {
dockerClient.execStartCmd(execCreateCmdResponse.getId()).exec(callback).awaitCompletion();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Python verification was interrupted.", e);
/** Closes the underlying Docker client, releasing its connections and threads. */
@Override
public void close() {
try {
// Unregister the shutdown hook so a closed executor is not retained by the JVM. Throws
// IllegalStateException if the JVM is already shutting down (e.g. close() invoked from the
// hook itself), in which case there is nothing to remove.
Runtime.getRuntime().removeShutdownHook(shutdownHook);
} catch (IllegalStateException e) {
// JVM shutdown already in progress; the hook cannot (and need not) be removed.
}
try {
dockerClient.close();
} catch (IOException e) {
throw new UncheckedIOException(e);
logger.warn("Failed to close docker client", e);
}
}

private void initContainer() {
if (dockerClient == null) {
throw new IllegalStateException("Docker client is not initialized.");
private static DockerClient buildDockerClient(String baseUrl) {
if (baseUrl != null) {
var config =
DefaultDockerClientConfig.createDefaultConfigBuilder().withDockerHost(baseUrl).build();
return DockerClientBuilder.getInstance(config).build();
}
return DockerClientBuilder.getInstance().build();
}

private void prepareImage() {
if (dockerPath != null) {
buildDockerImage();
} else {
Expand All @@ -203,34 +369,24 @@ private void initContainer() {
}
logger.info("Image {} is available.", image);
}
logger.info("Starting container for ContainerCodeExecutor...");
var createContainerResponse =
dockerClient.createContainerCmd(image).withTty(true).withAttachStdin(true).exec();
dockerClient.startContainerCmd(createContainerResponse.getId()).exec();

var containers = dockerClient.listContainersCmd().withShowAll(true).exec();
this.container =
containers.stream()
.filter(c -> c.getId().equals(createContainerResponse.getId()))
.findFirst()
.orElseThrow(() -> new IllegalStateException("Failed to find the created container."));

logger.info("Container {} started.", container.getId());
verifyPythonInstallation();
}

private void cleanupContainer() {
if (container == null) {
return;
private void buildDockerImage() {
if (dockerPath == null) {
throw new IllegalStateException("Docker path is not set.");
}
logger.info("[Cleanup] Stopping the container...");
dockerClient.stopContainerCmd(container.getId()).exec();
dockerClient.removeContainerCmd(container.getId()).exec();
logger.info("Container {} stopped and removed.", container.getId());
File dockerfile = new File(dockerPath);
if (!dockerfile.exists()) {
throw new UncheckedIOException(new IOException("Invalid Docker path: " + dockerPath));
}

logger.info("Building Docker image...");
try {
dockerClient.close();
} catch (IOException e) {
logger.warn("Failed to close docker client", e);
dockerClient.buildImageCmd(dockerfile).withTag(image).start().awaitCompletion();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException("Docker image build was interrupted.", e);
}
logger.info("Docker image: {} built.", image);
}
}
Loading
Loading