Skip to content

fix(dev): confine ReplayPlugin recordings to a configured replay root #1406

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+495 −39
Open

fix(dev): confine ReplayPlugin recordings to a configured replay root #1406

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
203 changes: 177 additions & 26 deletions dev/src/main/java/com/google/adk/plugins/ReplayPlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@
*/
package com.google.adk.plugins;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Strings.isNullOrEmpty;

import com.google.adk.agents.CallbackContext;
import com.google.adk.agents.InvocationContext;
import com.google.adk.models.LlmRequest;
Expand All @@ -27,38 +30,107 @@
import com.google.adk.tools.AgentTool;
import com.google.adk.tools.BaseTool;
import com.google.adk.tools.ToolContext;
import com.google.common.annotations.VisibleForTesting;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.InvalidPathException;
import java.nio.file.LinkOption;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Plugin for replaying ADK agent interactions from recordings. */
/**
* Plugin for replaying ADK agent interactions from recordings.
*
* <p>The replay case directory comes from the session state, which any caller of the dev server can
* set. Recordings are therefore only loaded from inside a replay root directory that the server
* operator configures: the {@code adk.replay.root} system property, the {@code ADK_REPLAY_ROOT}
* environment variable, or the process working directory when neither is set. Case directories that
* resolve outside the root, through {@code ..} segments, an absolute path or a symlink, are
* rejected. A relative case directory resolves against the root, not against the working directory.
*
* <p>The working-directory fallback rarely matches where the recordings live, so configure the root
* explicitly; the plugin logs the root it ended up with when it is constructed.
*/
public class ReplayPlugin extends BasePlugin {
private static final Logger logger = LoggerFactory.getLogger(ReplayPlugin.class);
private static final String REPLAY_CONFIG_KEY = "_adk_replay_config";
private static final String RECORDINGS_FILENAME = "generated-recordings.yaml";
private static final String REPLAY_ROOT_PROPERTY = "adk.replay.root";
private static final String REPLAY_ROOT_ENV = "ADK_REPLAY_ROOT";

// Track replay state per invocation to support concurrent runs
// key: invocation_id -> InvocationReplayState
private final Map<String, InvocationReplayState> invocationStates;

// Recordings are only read from inside this directory, never from arbitrary session state paths.
private final Path replayRoot;

public ReplayPlugin() {
this("adk_replay");
}

public ReplayPlugin(String name) {
this(name, defaultReplayRoot(), configuredReplayRoot() == null);
}

/** Creates a plugin that only loads recordings from inside {@code replayRoot}. */
@VisibleForTesting
ReplayPlugin(String name, Path replayRoot) {
this(name, replayRoot, /* usingWorkingDirectory= */ false);
}

private ReplayPlugin(String name, Path replayRoot, boolean usingWorkingDirectory) {
super(name);
this.replayRoot = checkNotNull(replayRoot).toAbsolutePath().normalize();
this.invocationStates = new ConcurrentHashMap<>();
logReplayRoot(usingWorkingDirectory);
}

private static Path defaultReplayRoot() {
String configured = configuredReplayRoot();
return Paths.get(configured != null ? configured : System.getProperty("user.dir", ""));
}

private static @Nullable String configuredReplayRoot() {
return configuredReplayRoot(
System.getProperty(REPLAY_ROOT_PROPERTY), System.getenv(REPLAY_ROOT_ENV));
}

/** Returns the operator-configured replay root, or null when neither knob carries a value. */
@VisibleForTesting
static @Nullable String configuredReplayRoot(
@Nullable String propertyValue, @Nullable String environmentValue) {
String configured = isNullOrEmpty(propertyValue) ? environmentValue : propertyValue;
return isNullOrEmpty(configured) ? null : configured;
}

/** Reports the effective replay root once, so a misconfigured root shows up at startup. */
private void logReplayRoot(boolean usingWorkingDirectory) {
String readable = Files.isDirectory(replayRoot) ? "" : " (not a readable directory)";
if (usingWorkingDirectory) {
logger.warn(
"Replay recordings are confined to the working directory {}{}, which is rarely where"
+ " recordings live. Set -D{} or {} to the directory holding them.",
replayRoot,
readable,
REPLAY_ROOT_PROPERTY,
REPLAY_ROOT_ENV);
} else {
logger.info("Replay recordings are confined to {}{}", replayRoot, readable);
}
}

@Override
Expand Down Expand Up @@ -169,17 +241,19 @@ private boolean isReplayModeOn(ToolContext toolContext) {
}

private boolean isReplayModeOnFromState(Map<String, Object> sessionState) {
if (!sessionState.containsKey(REPLAY_CONFIG_KEY)) {
return false;
}
Map<String, Object> config = replayConfig(sessionState);
return config != null && config.get("dir") != null && config.get("user_message_index") != null;
}

/** Returns the replay config from session state, or null when replay is not configured. */
private static Map<String, Object> replayConfig(Map<String, Object> sessionState) {
Object config = sessionState.get(REPLAY_CONFIG_KEY);
if (!(config instanceof Map)) {
return null;
}
@SuppressWarnings("unchecked")
Map<String, Object> config = (Map<String, Object>) sessionState.get(REPLAY_CONFIG_KEY);

String caseDir = (String) config.get("dir");
Integer msgIndex = (Integer) config.get("user_message_index");

return caseDir != null && msgIndex != null;
Map<String, Object> typedConfig = (Map<String, Object>) config;
return typedConfig;
}

private InvocationReplayState getInvocationState(CallbackContext callbackContext) {
Expand All @@ -194,44 +268,121 @@ private void loadInvocationState(InvocationContext invocationContext) {
String invocationId = invocationContext.invocationId();
Map<String, Object> sessionState = invocationContext.session().state();

@SuppressWarnings("unchecked")
Map<String, Object> config = (Map<String, Object>) sessionState.get(REPLAY_CONFIG_KEY);
Map<String, Object> config = replayConfig(sessionState);
if (config == null) {
throw new ReplayConfigError("Replay parameters are missing from session state");
}

String caseDir = (String) config.get("dir");
Integer msgIndex = (Integer) config.get("user_message_index");

if (caseDir == null || msgIndex == null) {
Object caseDirValue = config.get("dir");
Object msgIndexValue = config.get("user_message_index");
if (caseDirValue == null || msgIndexValue == null) {
throw new ReplayConfigError("Replay parameters are missing from session state");
}
if (!(caseDirValue instanceof String)) {
throw new ReplayConfigError(
"Replay parameter 'dir' must be a string, got "
+ caseDirValue.getClass().getSimpleName());
}
String caseDir = (String) caseDirValue;
int msgIndex = userMessageIndex(msgIndexValue);

// Load recordings
Path recordingsFile = Paths.get(caseDir, RECORDINGS_FILENAME);

if (!Files.exists(recordingsFile)) {
throw new ReplayConfigError("Recordings file not found: " + recordingsFile);
}
Path recordingsFile = resolveRecordingsFile(caseDir);

try {
Recordings recordings = RecordingsLoader.load(recordingsFile);
// NOFOLLOW_LINKS: the path was canonical when it was checked, so refuse it if it became a
// symlink in between.
try (InputStream recordingsStream =
Files.newInputStream(recordingsFile, LinkOption.NOFOLLOW_LINKS)) {
Recordings recordings = RecordingsLoader.load(recordingsStream);

// Create and store invocation state
InvocationReplayState state = new InvocationReplayState(caseDir, msgIndex, recordings);
invocationStates.put(invocationId, state);

// The case directory is caller input, so log the counts rather than the path.
logger.debug(
"Loaded replay state for invocation {}: case_dir={}, msg_index={}, recordings={}",
"Loaded replay state for invocation {}: msg_index={}, recordings={}",
invocationId,
caseDir,
msgIndex,
recordings.recordings().size());

} catch (IOException e) {
// The parser quotes the offending line, so neither the cause nor the file name is safe to
// carry: the name embeds the caller's case directory. Report the failure shape instead.
throw new ReplayConfigError(
"Failed to load the recordings file under "
+ replayRoot
+ " ("
+ e.getClass().getSimpleName()
+ ")");
}
}

/**
* Returns the user message index, which YAML and JSON decoders hand over as any numeric type.
*
* <p>Anything that is not a whole number in range is rejected rather than narrowed, so a bad
* index cannot quietly select the wrong recording.
*/
private static int userMessageIndex(Object value) {
if (!(value instanceof Number)) {
throw new ReplayConfigError(
"Failed to load recordings from " + recordingsFile + ": " + e.getMessage(), e);
"Replay parameter 'user_message_index' must be a number, got "
+ value.getClass().getSimpleName());
}
double asDouble = ((Number) value).doubleValue();
if (asDouble != Math.floor(asDouble) || asDouble < 0 || asDouble > Integer.MAX_VALUE) {
throw new ReplayConfigError(
"Replay parameter 'user_message_index' must be a whole number in [0, "
+ Integer.MAX_VALUE
+ "]");
}
return ((Number) value).intValue();
}

/**
* Resolves the recordings file for a session-supplied case directory, keeping it inside the
* replay root.
*
* <p>Containment is decided on the canonical path, so neither {@code ..} nor a symlink can leave
* the root however the caller spells it. That canonical path is what is returned, so the caller
* opens exactly what was checked. Messages never echo the case directory, which is caller input.
*/
private Path resolveRecordingsFile(String caseDir) {
Path canonicalRoot;
try {
canonicalRoot = replayRoot.toRealPath();
} catch (IOException e) {
throw new ReplayConfigError(
"Replay root directory is not readable: "
+ replayRoot
+ ". Set -D"
+ REPLAY_ROOT_PROPERTY
+ " to the directory holding the recordings.",
e);
}

Path recordingsFile;
try {
recordingsFile = replayRoot.resolve(caseDir).normalize().resolve(RECORDINGS_FILENAME);
} catch (InvalidPathException e) {
throw new ReplayConfigError("Replay parameter 'dir' is not a valid path", e);
}

Path canonicalFile;
try {
canonicalFile = recordingsFile.toRealPath();
} catch (NoSuchFileException e) {
throw new ReplayConfigError("Recordings file not found under the replay root " + replayRoot);
} catch (IOException e) {
throw new ReplayConfigError("Recordings file is not readable under " + replayRoot, e);
}
if (!canonicalFile.startsWith(canonicalRoot)) {
throw new ReplayConfigError(
"Replay directory resolves outside the replay root " + replayRoot);
}

return canonicalFile;
}

private Recording getNextRecordingForAgent(InvocationReplayState state, String agentName) {
Expand Down
Loading
Loading