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
108 changes: 108 additions & 0 deletions core/src/main/java/com/google/adk/plugins/PendingArtifactDelta.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
/*
* 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.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.google.adk.plugins;

import static com.google.common.collect.ImmutableMap.toImmutableMap;

import com.google.adk.agents.CallbackContext;
import com.google.adk.agents.InvocationContext;
import com.google.common.collect.ImmutableMap;
import java.util.Map;

/**
* Carries artifact versions from {@code onUserMessageCallback}, which has no {@link
* com.google.adk.events.EventActions} to write to, across to {@code beforeAgentCallback}, which
* does.
*
* <p>The hand-off rides on the session state under a {@link
* com.google.adk.sessions.State#TEMP_PREFIX} key, so {@code BaseSessionService.appendEvent} skips
* it when applying the state delta and the bookkeeping never reaches persisted session state. The
* key also carries the invocation id, so concurrent invocations sharing one session cannot read
* each other's pending versions.
*
* <p>Draining overwrites the entry with an empty map rather than removing it. {@code State.remove}
* would write the {@code State.REMOVED} sentinel into the event's state delta, which does not
* survive a JSON round trip.
*/
final class PendingArtifactDelta {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HIGH: The stash leaves one key in stored session state for every upload turn. Running two turns through InMemoryRunner left two temp:save_files_as_artifacts_plugin:pending_delta:<invocationId>={} entries. The direct write mutates the Session object stored by InMemorySessionService (L340-L358), and FirestoreSessionService persists it as well (L657-L696). Please follow Kotlin and keep the versions in per-invocation data (L107-L113, L120-L128), which in Java is InvocationContext.callbackContextData().


private static final String KEY = "temp:%s:pending_delta:%s";

private PendingArtifactDelta() {}

/** Stores the versions produced during this invocation. */
static void stash(
InvocationContext invocationContext, String pluginName, ImmutableMap<String, Integer> delta) {
invocationContext
.session()
.state()
.put(key(pluginName, invocationContext.invocationId()), delta);
}

/**
* Returns the stored versions and clears them, so only the first agent callback reports them.
*
* <p>The write is guarded on there being something to clear, and the guard is load-bearing: a
* write through {@link CallbackContext#state()} sets the state delta, and {@code BaseAgent} emits
* an event for any before-agent callback that left one. Clearing unconditionally would therefore
* emit an event carrying an empty artifact delta for every agent after the first.
*/
static ImmutableMap<String, Integer> drain(CallbackContext callbackContext, String pluginName) {
String key = key(pluginName, callbackContext.invocationId());
ImmutableMap<String, Integer> pending = read(callbackContext.state().get(key));
if (!pending.isEmpty()) {
callbackContext.state().put(key, ImmutableMap.of());
}
return pending;
}

/**
* Discards versions left behind by an invocation that ended before any agent ran, which happens
* when a {@code beforeRunCallback} on another plugin halts the run. There is no {@link
* com.google.adk.events.EventActions} to report them on at that point, so they are dropped rather
* than returned.
*
* <p>Only writes when something is actually stashed: an unconditional write would add an entry
* for every invocation that uploaded nothing, which is more state noise than the leak it
* prevents.
*/
static void clear(InvocationContext invocationContext, String pluginName) {
Map<String, Object> state = invocationContext.session().state();
String key = key(pluginName, invocationContext.invocationId());
if (!read(state.get(key)).isEmpty()) {
state.put(key, ImmutableMap.of());
}
}

private static String key(String pluginName, String invocationId) {
return KEY.formatted(pluginName, invocationId);
}

private static ImmutableMap<String, Integer> read(Object stashed) {
if (!(stashed instanceof Map<?, ?> entries)) {
return ImmutableMap.of();
}
return entries.entrySet().stream()
.filter(PendingArtifactDelta::isVersionEntry)
.collect(
toImmutableMap(entry -> (String) entry.getKey(), entry -> (Integer) entry.getValue()));
}

/** State values survive a JSON round trip untyped, so each entry is checked before it is kept. */
private static boolean isVersionEntry(Map.Entry<?, ?> entry) {
return entry.getKey() instanceof String && entry.getValue() instanceof Integer;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,227 @@
/*
* 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.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.google.adk.plugins;

import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;

import com.google.adk.agents.BaseAgent;
import com.google.adk.agents.CallbackContext;
import com.google.adk.agents.InvocationContext;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.List;
import java.util.Optional;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LOW: This Javadoc runs to about ten sentences, while Kotlin's covers the same points in three (L34-L45). Please trim it and use {@link com.google.adk.tools.LoadArtifactsTool} instead of naming load_artifacts.

* Plugin that saves files embedded in user messages as artifacts.
*
* <p>This allows users to upload files in the chat experience and have those files available to the
* agent within the current session. Each {@code inlineData} part of the incoming user message is
* written to the configured {@link com.google.adk.artifacts.BaseArtifactService} and replaced, in
* the message that reaches the model, by a short text placeholder naming the artifact. The bytes
* themselves are therefore stored once and not resent on every turn.
*
* <p>The artifact name is taken from {@link Blob#displayName()} when present, so an uploaded {@code
* report.pdf} is stored under that name. When the blob carries no display name, a name is generated
* from the invocation id and the part index. Uploading the same name again saves a new version of
* it, and the version saved is what this plugin reports for that name.
*
* <p>Add the {@code load_artifacts} tool to the agent, or load the artifacts from your own tool, to
* let the model read the stored bytes back.
*
* <p>Register it on the runner:
*
* <pre>{@code
* Runner runner =
* Runner.builder()
* .agent(agent)
* .appName("my-app")
* .artifactService(new InMemoryArtifactService())
* .sessionService(new InMemorySessionService())
* .plugins(new SaveFilesAsArtifactsPlugin())
* .build();
* }</pre>
*
* <p>The plugin is a no-op when no artifact service is configured on the runner.
*/
public class SaveFilesAsArtifactsPlugin extends BasePlugin {

/** Name used when the plugin is constructed without an explicit one. Matches adk-python's. */
public static final String DEFAULT_NAME = "save_files_as_artifacts_plugin";

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LOW: DEFAULT_NAME adds public API that neither Kotlin nor Python exposes, and no other plugin in this package exposes its default name. Please make it private, or package-private if the tests need it.


private static final Logger logger = LoggerFactory.getLogger(SaveFilesAsArtifactsPlugin.class);

private static final String GENERATED_FILE_NAME = "artifact_%s_%d";
private static final String PLACEHOLDER_TEXT = "[Uploaded Artifact: \"%s\"]";

public SaveFilesAsArtifactsPlugin() {
this(DEFAULT_NAME);
}

public SaveFilesAsArtifactsPlugin(String name) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HIGH: Please port attachFileReference as in Kotlin (L46-L51, L130-L171). No SPI change is needed. Kotlin reads the stored file's URI from saveAndReloadArtifact and its version from listVersions. Java's BaseArtifactService already provides both (L62-L73, L137-L142), and GcsArtifactService already returns a gs:// fileData from it (L266-L287). Because Python and Kotlin default it to true, adding it later would change what existing callers send to the model.

super(name);
}

@Override
public Maybe<Content> onUserMessageCallback(
InvocationContext invocationContext, Content userMessage) {
if (invocationContext.artifactService() == null) {
logger.warn("No artifact service is configured; plugin '{}' is disabled.", getName());
return Maybe.empty();
}
ImmutableList<Part> parts =
ImmutableList.copyOf(userMessage.parts().orElse(ImmutableList.of()));
if (parts.stream().noneMatch(SaveFilesAsArtifactsPlugin::hasInlineData)) {
return Maybe.empty();
}
return Flowable.range(0, parts.size())
.concatMapSingle(index -> savePart(invocationContext, parts.get(index), index))
.collect(toImmutableList())
.map(results -> rebuildMessage(invocationContext, userMessage, results))
.filter(Optional::isPresent)
.map(Optional::get);
}

/**
* Records the artifact versions stashed by {@link #onUserMessageCallback} on the first event
* actions of the invocation. {@code onUserMessageCallback} runs before any {@link
* com.google.adk.events.EventActions} exists, so the versions cannot be reported from there.
*
* <p>The reporting is a side effect and the return is always empty, deliberately: {@code
* PluginManager} stops at the first plugin that returns a value, so returning content here would
* both skip every later plugin's callback and halt the agent.
*/
@Override
public Maybe<Content> beforeAgentCallback(BaseAgent agent, CallbackContext callbackContext) {
PendingArtifactDelta.drain(callbackContext, getName())
.forEach(callbackContext.eventActions().artifactDelta()::put);
return Maybe.empty();
}

/**
* Discards a stash that {@link #beforeAgentCallback} never got to report, which happens when a
* {@code beforeRunCallback} on another plugin halts the invocation before any agent runs.
*
* <p>Only a run that completes reaches this hook; {@link #onRunErrorCallback} clears the same
* stash when one fails.
*/
@Override
public Completable afterRunCallback(InvocationContext invocationContext) {
PendingArtifactDelta.clear(invocationContext, getName());
return Completable.complete();
}

/**
* Discards the stash when the invocation fails, the window {@link #afterRunCallback} never sees.
*
* <p>The versions cannot be reported from here: this hook is handed an {@link InvocationContext},
* which carries no {@link com.google.adk.events.EventActions} to write an artifact delta to. The
* clear is idempotent, so a run reaching both hooks is no different from one reaching either.
*/
@Override
public Completable onRunErrorCallback(InvocationContext invocationContext, Throwable error) {
PendingArtifactDelta.clear(invocationContext, getName());
return Completable.complete();
}

/** Saves one part if it carries inline data, leaving every other part untouched. */
private Single<SavedPart> savePart(InvocationContext invocationContext, Part part, int index) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MEDIUM: Please port the 20 MB cap from Kotlin (L78-L89; Python L116-L128). An oversized file is not saved and is replaced with an [Upload Error: ...] note. Without the cap, the bytes can still reach the model here: LoadArtifactsTool, which the Javadoc recommends, puts a stored PDF or image back into the request inline (L225-L234), and a failed save also keeps the blob inline. Because an error note saves no version, return the rewritten message whenever a part changed, as Kotlin does.

if (!hasInlineData(part)) {
return Single.just(SavedPart.unchanged(part));
}
String fileName = resolveFileName(invocationContext, part, index);
return invocationContext
.artifactService()
.saveArtifact(
invocationContext.appName(),
invocationContext.userId(),
invocationContext.session().id(),
fileName,
part)
.map(version -> SavedPart.saved(placeholderFor(fileName), fileName, version))
.onErrorReturn(error -> keepOriginal(part, fileName, error));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HIGH: onErrorReturn only catches errors signaled by the returned Single. If the artifact service throws directly from saveArtifact, the whole turn fails. In a Runner test with such a service, the turn ended in IllegalStateException instead of keeping the original part. Kotlin catches any exception around the save (L90-L100). Wrapping the call in Single.defer(() -> ...saveArtifact(...)) fixes this.

}

/** A failed save must not fail the invocation: the original part is passed through unchanged. */
private SavedPart keepOriginal(Part part, String fileName, Throwable error) {
logger.error("Failed to save artifact '{}'; keeping the original part.", fileName, error);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HIGH: This logs the user-chosen file name, and the logged exception can repeat it. To avoid that, Kotlin logs only the part index and the exception type (L96-L97), for example logger.error("Failed to save the file in part {} ({}); keeping the original part.", index, error.getClass().getSimpleName()).

return SavedPart.unchanged(part);
}

/** Returns the rewritten message, or empty when no part was actually offloaded. */
private Optional<Content> rebuildMessage(
InvocationContext invocationContext, Content userMessage, List<SavedPart> results) {
ImmutableMap<String, Integer> delta = toArtifactDelta(results);
if (delta.isEmpty()) {
return Optional.empty();
}
PendingArtifactDelta.stash(invocationContext, getName(), delta);
ImmutableList<Part> parts = results.stream().map(SavedPart::part).collect(toImmutableList());
return Optional.of(userMessage.toBuilder().parts(parts).build());
}

private static ImmutableMap<String, Integer> toArtifactDelta(List<SavedPart> results) {
return results.stream()
.filter(SavedPart::isSaved)
.collect(
toImmutableMap(SavedPart::savedFileName, SavedPart::version, (older, newer) -> newer));
}

private static String resolveFileName(InvocationContext invocationContext, Part part, int index) {
return part.inlineData()
.flatMap(Blob::displayName)
.filter(displayName -> !displayName.isEmpty())
.orElseGet(() -> GENERATED_FILE_NAME.formatted(invocationContext.invocationId(), index));
}

private static Part placeholderFor(String fileName) {
return Part.fromText(PLACEHOLDER_TEXT.formatted(fileName));
}

private static boolean hasInlineData(Part part) {
return part.inlineData().isPresent();
}

/** One input part after the offload attempt: either untouched, or replaced by a placeholder. */
private record SavedPart(Part part, Optional<String> fileName, int version) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LOW: Kotlin handles this in one loop that builds the new parts and a map of versions (L66-L113). Using that structure would remove SavedPart (an Optional record component with a placeholder version), the single-use helpers, and the filter(Optional::isPresent).map(Optional::get) step. It would also keep the two ports easy to compare once the size cap and file reference are added.


static SavedPart unchanged(Part part) {
return new SavedPart(part, Optional.empty(), 0);
}

static SavedPart saved(Part placeholder, String fileName, int version) {
return new SavedPart(placeholder, Optional.of(fileName), version);
}

boolean isSaved() {
return fileName.isPresent();
}

String savedFileName() {
return fileName.orElseThrow();
}
}
}
Loading
Loading