From 272ea6cbcc8f61311e810e79d23d7f08dad31563 Mon Sep 17 00:00:00 2001 From: svetanis Date: Fri, 7 Aug 2026 18:00:43 -0700 Subject: [PATCH] feat(plugins): port SaveFilesAsArtifactsPlugin from adk-python --- .../java/com/google/adk/agents/BaseAgent.java | 3 +- .../plugins/SaveFilesAsArtifactsPlugin.java | 274 +++++++ .../com/google/adk/agents/BaseAgentTest.java | 140 ++++ .../SaveFilesAsArtifactsPluginTest.java | 735 ++++++++++++++++++ 4 files changed, 1151 insertions(+), 1 deletion(-) create mode 100644 core/src/main/java/com/google/adk/plugins/SaveFilesAsArtifactsPlugin.java create mode 100644 core/src/test/java/com/google/adk/plugins/SaveFilesAsArtifactsPluginTest.java diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index b14e5abe1..da43622b1 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -429,7 +429,8 @@ private Maybe callCallback( .switchIfEmpty( Maybe.defer( () -> { - if (callbackContext.state().hasDelta()) { + if (callbackContext.state().hasDelta() + || !callbackContext.eventActions().artifactDelta().isEmpty()) { Event.Builder eventBuilder = Event.builder() .id(invocationContext.newUuid()) diff --git a/core/src/main/java/com/google/adk/plugins/SaveFilesAsArtifactsPlugin.java b/core/src/main/java/com/google/adk/plugins/SaveFilesAsArtifactsPlugin.java new file mode 100644 index 000000000..220627d7e --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/SaveFilesAsArtifactsPlugin.java @@ -0,0 +1,274 @@ +/* + * 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 java.util.Comparator.naturalOrder; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.CallbackContext; +import com.google.adk.agents.InvocationContext; +import com.google.adk.artifacts.BaseArtifactService; +import com.google.adk.sessions.SessionKey; +import com.google.common.base.Ascii; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.genai.types.Blob; +import com.google.genai.types.Content; +import com.google.genai.types.FileData; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Optional; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Saves files attached to user messages as artifacts, replacing each with a placeholder that names + * the artifact so agents can load it later, such as with {@link + * com.google.adk.tools.LoadArtifactsTool}. Each file is saved under its {@link Blob#displayName() + * display name}, where a {@code user:} prefix shares it across the user's sessions, or under a + * generated name if it has no display name; files over 20 MB are replaced with an error note + * instead. Mirrors Python ADK's {@code SaveFilesAsArtifactsPlugin}. + */ +public class SaveFilesAsArtifactsPlugin extends BasePlugin { + + private static final String DEFAULT_NAME = "save_files_as_artifacts_plugin"; + + private static final Logger logger = LoggerFactory.getLogger(SaveFilesAsArtifactsPlugin.class); + + private static final int BYTES_PER_MB = 1024 * 1024; + + // The Gemini API's size limit for inline data. + private static final int MAX_FILE_SIZE_MB = 20; + private static final int MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * BYTES_PER_MB; + + // The URI schemes models can fetch: Vertex reads `gs`, hosted endpoints `https` and `http`. + private static final ImmutableSet MODEL_ACCESSIBLE_URI_SCHEMES = + ImmutableSet.of("gs", "https", "http"); + + private static final String PENDING_DELTA_KEY = "%s:pending_delta"; + private static final String GENERATED_FILE_NAME = "artifact_%s_%d"; + private static final String PLACEHOLDER = "[Uploaded Artifact: \"%s\"]"; + private static final String UPLOAD_ERROR = + "[Upload Error: File %s (%s MB) exceeds the maximum supported size of %dMB. Please upload a" + + " smaller file.]"; + + private final boolean attachFileReference; + private final String pendingDeltaKey; + + public SaveFilesAsArtifactsPlugin() { + this(/* attachFileReference= */ true); + } + + public SaveFilesAsArtifactsPlugin(boolean attachFileReference) { + this(attachFileReference, DEFAULT_NAME); + } + + /** + * Creates the plugin. + * + * @param attachFileReference Whether to also attach a {@link FileData} reference to each saved + * file when the artifact service stores it at a URI the model can read ({@code gs}, {@code + * https} or {@code http}). + * @param name The unique name of the plugin instance. + */ + public SaveFilesAsArtifactsPlugin(boolean attachFileReference, String name) { + super(name); + this.attachFileReference = attachFileReference; + this.pendingDeltaKey = PENDING_DELTA_KEY.formatted(name); + } + + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + BaseArtifactService artifactService = invocationContext.artifactService(); + if (artifactService == null) { + logger.warn("Artifact service is not set. SaveFilesAsArtifactsPlugin will not be enabled."); + return Maybe.empty(); + } + return Maybe.defer( + () -> { + ImmutableList parts = + ImmutableList.copyOf(userMessage.parts().orElse(ImmutableList.of())); + Map savedVersions = new LinkedHashMap<>(); + return Flowable.range(0, parts.size()) + .concatMapSingle( + index -> + newParts( + invocationContext, + artifactService, + parts.get(index), + index, + savedVersions)) + .concatMapIterable(replacements -> replacements) + .toList() + .filter(newParts -> !newParts.equals(parts)) + .map( + newParts -> { + if (!savedVersions.isEmpty()) { + // Unlike session state, which Python uses, per-invocation data leaves no + // entry behind. + invocationContext + .callbackContextData() + .put( + pendingDeltaKey, + new PendingDelta(ImmutableMap.copyOf(savedVersions))); + } + return userMessage.toBuilder().parts(ImmutableList.copyOf(newParts)).build(); + }); + }); + } + + /** + * Reports the versions saved from the user message on the event of the first agent to run, since + * a plugin cannot add actions to the user's own event. + */ + @Override + public Maybe beforeAgentCallback(BaseAgent agent, CallbackContext callbackContext) { + return Maybe.fromAction( + () -> { + if (callbackContext.invocationContext().callbackContextData().remove(pendingDeltaKey) + instanceof PendingDelta pending) { + callbackContext.eventActions().artifactDelta().putAll(pending.versions()); + } + }); + } + + /** Returns the parts that take {@code part}'s place, recording any version saved. */ + private Single> newParts( + InvocationContext invocationContext, + BaseArtifactService artifactService, + Part part, + int index, + Map savedVersions) { + Optional inlineData = part.inlineData(); + if (inlineData.isEmpty()) { + return Single.just(ImmutableList.of(part)); + } + String fileName = + inlineData + .flatMap(Blob::displayName) + .filter(displayName -> !displayName.isEmpty()) + .orElseGet( + () -> GENERATED_FILE_NAME.formatted(invocationContext.invocationId(), index)); + int size = inlineData.flatMap(Blob::data).map(data -> data.length).orElse(0); + if (size > MAX_FILE_SIZE_BYTES) { + logger.warn( + "The file in part {} is {} bytes, over the limit; it was not saved.", index, size); + return Single.just( + ImmutableList.of( + Part.fromText(UPLOAD_ERROR.formatted(fileName, megabytes(size), MAX_FILE_SIZE_MB)))); + } + SessionKey sessionKey = invocationContext.session().sessionKey(); + return Single.defer(() -> save(artifactService, sessionKey, fileName, part)) + .map( + saved -> { + saved.version().ifPresent(version -> savedVersions.put(fileName, version)); + ImmutableList.Builder replacement = ImmutableList.builder(); + replacement.add(Part.fromText(PLACEHOLDER.formatted(fileName))); + saved.reference().ifPresent(replacement::add); + return replacement.build(); + }) + .onErrorReturn( + error -> { + // The exception message can contain the user-chosen file name, so log only the type. + logger.error( + "Failed to save the file in part {} ({}); keeping the original part.", + index, + error.getClass().getSimpleName()); + return ImmutableList.of(part); + }); + } + + /** + * Saves {@code part} under {@code fileName}. When {@code attachFileReference} is enabled, this + * uses {@link BaseArtifactService#saveAndReloadArtifact} to get the stored file's URI and then + * reads back the latest version, since the service cannot look up the URI of a given version. + */ + private Single save( + BaseArtifactService artifactService, SessionKey sessionKey, String fileName, Part part) { + if (!attachFileReference) { + return artifactService + .saveArtifact(sessionKey, fileName, part) + .map(version -> new SavedFile(Optional.of(version), Optional.empty())); + } + return artifactService + .saveAndReloadArtifact(sessionKey, fileName, part) + .flatMap( + stored -> { + Optional storedFile = stored.fileData(); + Optional reference = + storedFile + .flatMap(FileData::fileUri) + .filter(SaveFilesAsArtifactsPlugin::isModelAccessibleUri) + .map( + uri -> { + FileData.Builder fileData = + FileData.builder().displayName(fileName).fileUri(uri); + part.inlineData() + .flatMap(Blob::mimeType) + .filter(type -> !type.isEmpty()) + .or(() -> storedFile.flatMap(FileData::mimeType)) + .ifPresent(fileData::mimeType); + return Part.builder().fileData(fileData.build()).build(); + }); + // The file is saved by now, so failing to read its version only leaves it out of + // the delta. + return Single.defer(() -> artifactService.listVersions(sessionKey, fileName)) + .map(versions -> versions.stream().max(naturalOrder())) + .doOnSuccess( + version -> { + if (version.isEmpty()) { + logger.warn("Saved a file but the service listed no versions of it."); + } + }) + .onErrorReturn( + error -> { + logger.warn( + "Saved a file but could not list its versions ({}).", + error.getClass().getSimpleName()); + return Optional.empty(); + }) + .map(version -> new SavedFile(version, reference)); + }); + } + + private static boolean isModelAccessibleUri(String uri) { + int colon = uri.indexOf(':'); + return colon > 0 + && MODEL_ACCESSIBLE_URI_SCHEMES.contains(Ascii.toLowerCase(uri.substring(0, colon))); + } + + /** + * Formats {@code bytes} in megabytes to two decimal places, rounding half to even like Python. + */ + private static String megabytes(int bytes) { + return BigDecimal.valueOf(bytes) + .divide(BigDecimal.valueOf(BYTES_PER_MB)) + .setScale(2, RoundingMode.HALF_EVEN) + .toPlainString(); + } + + private record SavedFile(Optional version, Optional reference) {} + + private record PendingDelta(ImmutableMap versions) {} +} diff --git a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java index 4d803edbe..2c0b98d3e 100644 --- a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java @@ -381,6 +381,76 @@ public Maybe call(CallbackContext context) { assertThat(beforeCallback.wasCalled()).isTrue(); } + @Test + public void + runAsync_beforeCallbackSavesArtifactAndReturnsEmpty_invokesRunAsyncImplAndReturnsArtifactEvent() { + var runAsyncImpl = TestCallback.returningEmpty(); + Content runAsyncImplContent = Content.fromParts(Part.fromText("main_output")); + BeforeAgentCallback beforeCallback = + new BeforeAgentCallback() { + @Override + public Maybe call(CallbackContext context) { + return context.saveArtifact("file", Part.fromText("contents")).andThen(Maybe.empty()); + } + }; + var afterCallback = TestCallback.returningEmpty(); + TestBaseAgent agent = + new TestBaseAgent( + TEST_AGENT_NAME, + TEST_AGENT_DESCRIPTION, + ImmutableList.of(beforeCallback), + ImmutableList.of(afterCallback.asAfterAgentCallback()), + runAsyncImpl.asRunAsyncImplSupplier(runAsyncImplContent)); + InvocationContext invocationContext = TestUtils.createInvocationContext(agent); + + List results = agent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(results).hasSize(2); + // Artifact event from before callback + assertThat(results.get(0).content()).isEmpty(); + assertThat(results.get(0).actions().artifactDelta()).containsExactly("file", 0); + assertThat(results.get(0).actions().stateDelta()).isEmpty(); + // Content event from runAsyncImpl + assertThat(results.get(1).content()).hasValue(runAsyncImplContent); + assertThat(runAsyncImpl.wasCalled()).isTrue(); + assertThat(afterCallback.wasCalled()).isTrue(); + } + + @Test + public void + runAsync_afterCallbackSavesArtifactAndReturnsEmpty_invokesRunAsyncImplAndReturnsArtifactEvent() { + var runAsyncImpl = TestCallback.returningEmpty(); + Content runAsyncImplContent = Content.fromParts(Part.fromText("main_output")); + var beforeCallback = TestCallback.returningEmpty(); + AfterAgentCallback afterCallback = + new AfterAgentCallback() { + @Override + public Maybe call(CallbackContext context) { + return context.saveArtifact("file", Part.fromText("contents")).andThen(Maybe.empty()); + } + }; + TestBaseAgent agent = + new TestBaseAgent( + TEST_AGENT_NAME, + TEST_AGENT_DESCRIPTION, + ImmutableList.of(beforeCallback.asBeforeAgentCallback()), + ImmutableList.of(afterCallback), + runAsyncImpl.asRunAsyncImplSupplier(runAsyncImplContent)); + InvocationContext invocationContext = TestUtils.createInvocationContext(agent); + + List results = agent.runAsync(invocationContext).toList().blockingGet(); + + assertThat(results).hasSize(2); + // Content event from runAsyncImpl + assertThat(results.get(0).content()).hasValue(runAsyncImplContent); + // Artifact event from after callback + assertThat(results.get(1).content()).isEmpty(); + assertThat(results.get(1).actions().artifactDelta()).containsExactly("file", 0); + assertThat(results.get(1).actions().stateDelta()).isEmpty(); + assertThat(runAsyncImpl.wasCalled()).isTrue(); + assertThat(beforeCallback.wasCalled()).isTrue(); + } + @Test public void runAsync_firstAfterCallbackReturnsContent_skipsSecondAfterCallback() { var runAsyncImpl = TestCallback.returningEmpty(); @@ -654,6 +724,76 @@ public Maybe call(CallbackContext context) { assertThat(beforeCallback.wasCalled()).isTrue(); } + @Test + public void + runLive_beforeCallbackSavesArtifactAndReturnsEmpty_invokesRunLiveImplAndReturnsArtifactEvent() { + var runLiveImpl = TestCallback.returningEmpty(); + Content runLiveImplContent = Content.fromParts(Part.fromText("main_output")); + BeforeAgentCallback beforeCallback = + new BeforeAgentCallback() { + @Override + public Maybe call(CallbackContext context) { + return context.saveArtifact("file", Part.fromText("contents")).andThen(Maybe.empty()); + } + }; + var afterCallback = TestCallback.returningEmpty(); + TestBaseAgent agent = + new TestBaseAgent( + TEST_AGENT_NAME, + TEST_AGENT_DESCRIPTION, + ImmutableList.of(beforeCallback), + ImmutableList.of(afterCallback.asAfterAgentCallback()), + runLiveImpl.asRunLiveImplSupplier(runLiveImplContent)); + InvocationContext invocationContext = TestUtils.createInvocationContext(agent); + + List results = agent.runLive(invocationContext).toList().blockingGet(); + + assertThat(results).hasSize(2); + // Artifact event from before callback + assertThat(results.get(0).content()).isEmpty(); + assertThat(results.get(0).actions().artifactDelta()).containsExactly("file", 0); + assertThat(results.get(0).actions().stateDelta()).isEmpty(); + // Content event from runLiveImpl + assertThat(results.get(1).content()).hasValue(runLiveImplContent); + assertThat(runLiveImpl.wasCalled()).isTrue(); + assertThat(afterCallback.wasCalled()).isTrue(); + } + + @Test + public void + runLive_afterCallbackSavesArtifactAndReturnsEmpty_invokesRunLiveImplAndReturnsArtifactEvent() { + var runLiveImpl = TestCallback.returningEmpty(); + Content runLiveImplContent = Content.fromParts(Part.fromText("main_output")); + var beforeCallback = TestCallback.returningEmpty(); + AfterAgentCallback afterCallback = + new AfterAgentCallback() { + @Override + public Maybe call(CallbackContext context) { + return context.saveArtifact("file", Part.fromText("contents")).andThen(Maybe.empty()); + } + }; + TestBaseAgent agent = + new TestBaseAgent( + TEST_AGENT_NAME, + TEST_AGENT_DESCRIPTION, + ImmutableList.of(beforeCallback.asBeforeAgentCallback()), + ImmutableList.of(afterCallback), + runLiveImpl.asRunLiveImplSupplier(runLiveImplContent)); + InvocationContext invocationContext = TestUtils.createInvocationContext(agent); + + List results = agent.runLive(invocationContext).toList().blockingGet(); + + assertThat(results).hasSize(2); + // Content event from runLiveImpl + assertThat(results.get(0).content()).hasValue(runLiveImplContent); + // Artifact event from after callback + assertThat(results.get(1).content()).isEmpty(); + assertThat(results.get(1).actions().artifactDelta()).containsExactly("file", 0); + assertThat(results.get(1).actions().stateDelta()).isEmpty(); + assertThat(runLiveImpl.wasCalled()).isTrue(); + assertThat(beforeCallback.wasCalled()).isTrue(); + } + @Test public void runLive_firstAfterCallbackReturnsContent_skipsSecondAfterCallback() { var runLiveImpl = TestCallback.returningEmpty(); diff --git a/core/src/test/java/com/google/adk/plugins/SaveFilesAsArtifactsPluginTest.java b/core/src/test/java/com/google/adk/plugins/SaveFilesAsArtifactsPluginTest.java new file mode 100644 index 000000000..85ff80aca --- /dev/null +++ b/core/src/test/java/com/google/adk/plugins/SaveFilesAsArtifactsPluginTest.java @@ -0,0 +1,735 @@ +/* + * 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.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createRootAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static java.nio.charset.StandardCharsets.UTF_8; + +import com.google.adk.agents.CallbackContext; +import com.google.adk.agents.InvocationContext; +import com.google.adk.artifacts.BaseArtifactService; +import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.artifacts.ListArtifactsResponse; +import com.google.adk.events.Event; +import com.google.adk.runner.InMemoryRunner; +import com.google.adk.sessions.Session; +import com.google.adk.sessions.SessionKey; +import com.google.adk.testing.TestLlm; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Iterables; +import com.google.genai.types.Blob; +import com.google.genai.types.Content; +import com.google.genai.types.FileData; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class SaveFilesAsArtifactsPluginTest { + + private static final int MAX_SIZE = 20 * 1024 * 1024; + private static final String INVOCATION_ID = "test-invocation-id"; + private static final String MIME_TYPE = "application/pdf"; + private static final String PLACEHOLDER = "[Uploaded Artifact: \"%s\"]"; + private static final String UPLOAD_ERROR = + "[Upload Error: File large.pdf (%s MB) exceeds the maximum supported size of 20MB. Please" + + " upload a smaller file.]"; + private static final String STORED_URI = "%s%s/%d"; + private static final String GS_PREFIX = "gs://bucket/"; + private static final String APP_NAME = "InMemoryRunner"; + private static final String USER_ID = "user1"; + + private final FakeArtifactService service = FakeArtifactService.create(); + private final InvocationContext context = invocationContext(service); + private final SessionKey sessionKey = context.session().sessionKey(); + + @Test + public void constructor_defaultName() { + assertThat(new SaveFilesAsArtifactsPlugin().getName()) + .isEqualTo("save_files_as_artifacts_plugin"); + } + + @Test + public void onUserMessage_savesTheFileAndPutsAPlaceholderAndReferenceInItsPlace() { + Part file = filePart("report.pdf"); + + Optional result = + onUserMessage(new SaveFilesAsArtifactsPlugin(), context, userMessage(file)); + + assertThat(service.saves).containsExactly(new Save("report.pdf", file)); + assertThat(result) + .hasValue( + userMessage( + placeholder("report.pdf"), reference("report.pdf", "gs://bucket/report.pdf/0"))); + } + + @Test + public void onUserMessage_attachFileReferenceFalse_putsOnlyAPlaceholderAndReportsTheVersion() { + SaveFilesAsArtifactsPlugin plugin = + new SaveFilesAsArtifactsPlugin(/* attachFileReference= */ false); + int earlierVersion = + service + .saveArtifact(sessionKey, "report.pdf", filePart("report.pdf", "old contents")) + .blockingGet(); + Part file = filePart("report.pdf"); + + Optional result = onUserMessage(plugin, context, userMessage(file)); + + assertThat(Iterables.getLast(service.saves)).isEqualTo(new Save("report.pdf", file)); + assertThat(result).hasValue(userMessage(placeholder("report.pdf"))); + assertThat(beforeAgent(plugin, context)).containsExactly("report.pdf", earlierVersion + 1); + } + + @Test + public void onUserMessage_uriTheModelCannotRead_putsOnlyAPlaceholder() { + InvocationContext fileContext = + invocationContext(FakeArtifactService.withUriPrefix("file:///tmp/")); + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), fileContext, userMessage(filePart("report.pdf"))); + + assertThat(result).hasValue(userMessage(placeholder("report.pdf"))); + } + + @Test + public void onUserMessage_uriWithoutScheme_putsOnlyAPlaceholder() { + InvocationContext noSchemeContext = + invocationContext(FakeArtifactService.withUriPrefix("bucket/")); + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), noSchemeContext, userMessage(filePart("report.pdf"))); + + assertThat(result).hasValue(userMessage(placeholder("report.pdf"))); + } + + @Test + public void onUserMessage_httpsUri_addsTheReference() { + InvocationContext httpsContext = + invocationContext(FakeArtifactService.withUriPrefix("https://example.com/")); + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), httpsContext, userMessage(filePart("report.pdf"))); + + assertThat(result) + .hasValue( + userMessage( + placeholder("report.pdf"), + reference("report.pdf", "https://example.com/report.pdf/0"))); + } + + @Test + public void onUserMessage_upperCaseUriScheme_addsTheReference() { + InvocationContext upperCaseContext = + invocationContext(FakeArtifactService.withUriPrefix("HTTP://example.com/")); + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), + upperCaseContext, + userMessage(filePart("report.pdf"))); + + assertThat(result) + .hasValue( + userMessage( + placeholder("report.pdf"), + reference("report.pdf", "HTTP://example.com/report.pdf/0"))); + } + + @Test + public void onUserMessage_fileWithoutMimeType_takesTheReferenceTypeFromTheStoredFile() { + Part file = + Part.builder() + .inlineData( + Blob.builder().mimeType("").displayName("notes").data("text".getBytes(UTF_8))) + .build(); + + Optional result = + onUserMessage(new SaveFilesAsArtifactsPlugin(), context, userMessage(file)); + + Part reference = + Part.builder() + .fileData( + FileData.builder() + .mimeType("text/plain") + .displayName("notes") + .fileUri("gs://bucket/notes/0")) + .build(); + assertThat(result).hasValue(userMessage(placeholder("notes"), reference)); + } + + @Test + public void onUserMessage_inMemoryService_putsOnlyAPlaceholder() { + InMemoryArtifactService inMemory = new InMemoryArtifactService(); + InvocationContext inMemoryContext = invocationContext(inMemory); + Part file = filePart("report.pdf"); + + Optional result = + onUserMessage(new SaveFilesAsArtifactsPlugin(), inMemoryContext, userMessage(file)); + + assertThat(result).hasValue(userMessage(placeholder("report.pdf"))); + SessionKey inMemoryKey = inMemoryContext.session().sessionKey(); + assertThat(inMemory.loadArtifact(inMemoryKey, "report.pdf").blockingGet()).isEqualTo(file); + } + + @Test + public void onUserMessage_noDisplayName_savesUnderAGeneratedName() { + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(/* attachFileReference= */ false), + context, + userMessage(Part.fromText("See:"), filePart(""))); + + String generatedName = "artifact_test-invocation-id_1"; + assertThat(service.savedNames()).containsExactly(generatedName); + assertThat(result).hasValue(userMessage(Part.fromText("See:"), placeholder(generatedName))); + } + + @Test + public void onUserMessage_blobWithoutDisplayName_savesUnderAGeneratedName() { + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(/* attachFileReference= */ false), + context, + userMessage(Part.fromText("See:"), Part.fromBytes("text".getBytes(UTF_8), MIME_TYPE))); + + String generatedName = "artifact_test-invocation-id_1"; + assertThat(service.savedNames()).containsExactly(generatedName); + assertThat(result).hasValue(userMessage(Part.fromText("See:"), placeholder(generatedName))); + } + + @Test + public void onUserMessage_severalFiles_keepsTheOtherPartsInPlace() { + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), + context, + userMessage( + filePart("a.txt"), Part.fromText("Some text between files"), filePart("b.jpg"))); + + assertThat(result) + .hasValue( + userMessage( + placeholder("a.txt"), + reference("a.txt", "gs://bucket/a.txt/0"), + Part.fromText("Some text between files"), + placeholder("b.jpg"), + reference("b.jpg", "gs://bucket/b.jpg/0"))); + } + + @Test + public void onUserMessage_sameNameTwice_savesBothAndReportsTheLatestVersion() { + SaveFilesAsArtifactsPlugin plugin = new SaveFilesAsArtifactsPlugin(); + + Optional result = + onUserMessage( + plugin, + context, + userMessage(filePart("report.pdf", "first"), filePart("report.pdf", "second"))); + + assertThat(result) + .hasValue( + userMessage( + placeholder("report.pdf"), + reference("report.pdf", "gs://bucket/report.pdf/0"), + placeholder("report.pdf"), + reference("report.pdf", "gs://bucket/report.pdf/1"))); + assertThat(beforeAgent(plugin, context)).containsExactly("report.pdf", 1); + } + + @Test + public void onUserMessage_noArtifactService_returnsTheMessageUnchanged() { + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), + invocationContext(/* artifactService= */ null), + userMessage(filePart("report.pdf"))); + + assertThat(result).isEmpty(); + } + + @Test + public void onUserMessage_noFiles_returnsTheMessageUnchanged() { + Content message = userMessage(Part.fromText("Hello"), Part.fromText("No files here")); + + assertThat(onUserMessage(new SaveFilesAsArtifactsPlugin(), context, message)).isEmpty(); + assertThat(service.saves).isEmpty(); + } + + @Test + public void onUserMessage_noParts_returnsTheMessageUnchanged() { + Content message = Content.builder().role("user").build(); + + assertThat(onUserMessage(new SaveFilesAsArtifactsPlugin(), context, message)).isEmpty(); + assertThat(service.saves).isEmpty(); + } + + @Test + public void onUserMessage_savesNothingUntilSubscribed() { + Maybe pending = + new SaveFilesAsArtifactsPlugin() + .onUserMessageCallback(context, userMessage(filePart("report.pdf"))); + + assertThat(service.saves).isEmpty(); + pending.test().assertComplete(); + assertThat(service.savedNames()).containsExactly("report.pdf"); + } + + @Test + public void onUserMessage_saveFails_keepsTheFile() { + InvocationContext failingContext = + invocationContext(FakeArtifactService.failingOn("report.pdf")); + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), failingContext, userMessage(filePart("report.pdf"))); + + assertThat(result).isEmpty(); + } + + @Test + public void onUserMessage_saveReturnsAnError_keepsTheFile() { + InvocationContext erroringContext = + invocationContext(FakeArtifactService.erroringOn("report.pdf")); + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), erroringContext, userMessage(filePart("report.pdf"))); + + assertThat(result).isEmpty(); + } + + @Test + public void onUserMessage_oneOfTwoSavesFails_keepsOnlyTheFileThatFailed() { + InvocationContext failingContext = + invocationContext(FakeArtifactService.failingOn("failure.pdf")); + Part failed = filePart("failure.pdf"); + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), + failingContext, + userMessage(filePart("success.pdf"), failed)); + + assertThat(result) + .hasValue( + userMessage( + placeholder("success.pdf"), + reference("success.pdf", "gs://bucket/success.pdf/0"), + failed)); + } + + @Test + public void onUserMessage_versionCannotBeReadBack_keepsThePlaceholderAndReportsNoVersion() { + SaveFilesAsArtifactsPlugin plugin = new SaveFilesAsArtifactsPlugin(); + InvocationContext unreadable = invocationContext(FakeArtifactService.withUnreadableVersions()); + + Optional result = + onUserMessage(plugin, unreadable, userMessage(filePart("report.pdf"))); + + assertThat(result) + .hasValue( + userMessage( + placeholder("report.pdf"), reference("report.pdf", "gs://bucket/report.pdf/0"))); + assertThat(beforeAgent(plugin, unreadable)).isEmpty(); + } + + @Test + public void onUserMessage_serviceListsNoVersions_reportsNoVersion() { + SaveFilesAsArtifactsPlugin plugin = new SaveFilesAsArtifactsPlugin(); + InvocationContext dummy = invocationContext(new DummyArtifactService()); + + Optional result = onUserMessage(plugin, dummy, userMessage(filePart("report.pdf"))); + + assertThat(result).hasValue(userMessage(placeholder("report.pdf"))); + assertThat(beforeAgent(plugin, dummy)).isEmpty(); + } + + @Test + public void onUserMessage_fileAtTheSizeLimit_isSaved() { + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(/* attachFileReference= */ false), + context, + userMessage(filePart("max.pdf", MAX_SIZE))); + + assertThat(service.savedNames()).containsExactly("max.pdf"); + assertThat(result).hasValue(userMessage(placeholder("max.pdf"))); + } + + @Test + public void onUserMessage_fileOverTheSizeLimit_isReplacedWithAnError() { + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), + context, + userMessage(filePart("small.pdf"), filePart("large.pdf", MAX_SIZE + 1))); + + assertThat(service.savedNames()).containsExactly("small.pdf"); + assertThat(result) + .hasValue( + userMessage( + placeholder("small.pdf"), + reference("small.pdf", "gs://bucket/small.pdf/0"), + Part.fromText(UPLOAD_ERROR.formatted("20.00")))); + } + + @Test + public void onUserMessage_fileSizeHalfwayBetweenHundredths_roundsHalfToEven() { + // 20.125 MB, which Python's `:.2f` formats as 20.12. + int size = MAX_SIZE + MAX_SIZE / 160; + + Optional result = + onUserMessage( + new SaveFilesAsArtifactsPlugin(), context, userMessage(filePart("large.pdf", size))); + + assertThat(result).hasValue(userMessage(Part.fromText(UPLOAD_ERROR.formatted("20.12")))); + } + + @Test + public void beforeAgent_reportsTheSavedVersionsOnce() { + SaveFilesAsArtifactsPlugin plugin = new SaveFilesAsArtifactsPlugin(); + int earlierVersion = service.saveArtifact(sessionKey, "a.txt", filePart("a.txt")).blockingGet(); + assertThat(onUserMessage(plugin, context, userMessage(filePart("a.txt"), filePart("b.txt")))) + .isPresent(); + + ImmutableMap first = beforeAgent(plugin, context); + ImmutableMap second = beforeAgent(plugin, context); + + assertThat(first).containsExactly("a.txt", earlierVersion + 1, "b.txt", 0); + assertThat(second).isEmpty(); + } + + @Test + public void beforeAgent_oneOfTwoSavesFails_reportsOnlyTheSavedVersion() { + SaveFilesAsArtifactsPlugin plugin = new SaveFilesAsArtifactsPlugin(); + InvocationContext failingContext = + invocationContext(FakeArtifactService.failingOn("failure.pdf")); + assertThat( + onUserMessage( + plugin, + failingContext, + userMessage(filePart("success.pdf"), filePart("failure.pdf")))) + .isPresent(); + + assertThat(beforeAgent(plugin, failingContext)).containsExactly("success.pdf", 0); + } + + @Test + public void beforeAgent_twoPluginNames_eachReportsOnlyItsOwnVersions() { + SaveFilesAsArtifactsPlugin first = + new SaveFilesAsArtifactsPlugin(/* attachFileReference= */ true, "first_plugin"); + SaveFilesAsArtifactsPlugin second = + new SaveFilesAsArtifactsPlugin(/* attachFileReference= */ true, "second_plugin"); + assertThat(onUserMessage(first, context, userMessage(filePart("a.txt")))).isPresent(); + assertThat(onUserMessage(second, context, userMessage(filePart("b.txt")))).isPresent(); + + assertThat(beforeAgent(first, context)).containsExactly("a.txt", 0); + assertThat(beforeAgent(second, context)).containsExactly("b.txt", 0); + } + + @Test + public void beforeAgent_nothingSaved_reportsNothing() { + assertThat(beforeAgent(new SaveFilesAsArtifactsPlugin(), context)).isEmpty(); + } + + @Test + public void beforeAgent_reportsNothingUntilSubscribed() { + SaveFilesAsArtifactsPlugin plugin = new SaveFilesAsArtifactsPlugin(); + assertThat(onUserMessage(plugin, context, userMessage(filePart("report.pdf")))).isPresent(); + CallbackContext callbackContext = new CallbackContext(context, /* eventActions= */ null); + + Maybe pending = plugin.beforeAgentCallback(context.agent(), callbackContext); + + assertThat(callbackContext.eventActions().artifactDelta()).isEmpty(); + pending.test().assertComplete().assertNoValues(); + assertThat(callbackContext.eventActions().artifactDelta()).containsExactly("report.pdf", 0); + } + + @Test + public void runAsync_modelGetsThePlaceholderAndTheAgentEventReportsTheVersion() { + TestLlm model = + createTestLlm(createTextLlmResponse("Got it."), createTextLlmResponse("Got it.")); + InMemoryRunner runner = + new InMemoryRunner( + createTestAgentBuilder(model).build(), + APP_NAME, + ImmutableList.of(new SaveFilesAsArtifactsPlugin())); + Session session = runner.sessionService().createSession(APP_NAME, USER_ID).blockingGet(); + + ImmutableList> turns = + ImmutableList.of("first", "second").stream() + .map(text -> runTurn(runner, session, text)) + .collect(toImmutableList()); + + Content expectedUserContent = userMessage(Part.fromText("Read this"), placeholder("notes.txt")); + assertThat(Iterables.getLast(Iterables.getLast(model.getRequests()).contents())) + .isEqualTo(expectedUserContent); + assertThat( + turns.stream() + .map(events -> events.get(0).actions().artifactDelta()) + .collect(toImmutableList())) + .containsExactly(ImmutableMap.of("notes.txt", 0), ImmutableMap.of("notes.txt", 1)) + .inOrder(); + Session stored = + runner + .sessionService() + .getSession(APP_NAME, USER_ID, session.id(), Optional.empty()) + .blockingGet(); + assertThat(stored.immutableEvents().get(0).content()).hasValue(expectedUserContent); + assertThat(stored.state()).isEmpty(); + Part storedFile = + runner.artifactService().loadArtifact(stored.sessionKey(), "notes.txt").blockingGet(); + assertThat(new String(storedFile.inlineData().flatMap(Blob::data).orElseThrow(), UTF_8)) + .isEqualTo("second"); + } + + private static ImmutableList runTurn(InMemoryRunner runner, Session session, String text) { + Content message = userMessage(Part.fromText("Read this"), filePart("notes.txt", text)); + return ImmutableList.copyOf( + runner.runAsync(USER_ID, session.id(), message).toList().blockingGet()); + } + + private static InvocationContext invocationContext( + @Nullable BaseArtifactService artifactService) { + return createInvocationContext(createRootAgent()).toBuilder() + .artifactService(artifactService) + .invocationId(INVOCATION_ID) + .build(); + } + + /** The rewritten message, or empty when the plugin leaves the message unchanged. */ + private static Optional onUserMessage( + SaveFilesAsArtifactsPlugin plugin, InvocationContext invocationContext, Content message) { + return Optional.ofNullable( + plugin.onUserMessageCallback(invocationContext, message).blockingGet()); + } + + /** The artifact delta the plugin reports on the first agent's event. */ + private static ImmutableMap beforeAgent( + SaveFilesAsArtifactsPlugin plugin, InvocationContext invocationContext) { + CallbackContext callbackContext = + new CallbackContext(invocationContext, /* eventActions= */ null); + plugin + .beforeAgentCallback(invocationContext.agent(), callbackContext) + .test() + .assertComplete() + .assertNoValues(); + return ImmutableMap.copyOf(callbackContext.eventActions().artifactDelta()); + } + + private static Content userMessage(Part... parts) { + return Content.builder().role("user").parts(parts).build(); + } + + private static Part filePart(String displayName) { + return filePart(displayName, "file contents"); + } + + private static Part filePart(String displayName, String text) { + return filePart(displayName, text.getBytes(UTF_8)); + } + + private static Part filePart(String displayName, int size) { + return filePart(displayName, new byte[size]); + } + + private static Part filePart(String displayName, byte[] data) { + return Part.builder() + .inlineData(Blob.builder().mimeType(MIME_TYPE).displayName(displayName).data(data)) + .build(); + } + + private static Part placeholder(String fileName) { + return Part.fromText(PLACEHOLDER.formatted(fileName)); + } + + private static Part reference(String fileName, String uri) { + return Part.builder() + .fileData(FileData.builder().mimeType(MIME_TYPE).displayName(fileName).fileUri(uri)) + .build(); + } + + private record Save(String fileName, Part artifact) {} + + /** + * Stores artifacts in memory and reports each saved version at {@code + * /}. Saving a name in {@code failingNames} throws, saving one in + * {@code erroringNames} returns a failed {@link Single}, and listing versions throws unless + * {@code versionsReadable}. + */ + private static final class FakeArtifactService implements BaseArtifactService { + private final String uriPrefix; + private final ImmutableSet failingNames; + private final ImmutableSet erroringNames; + private final boolean versionsReadable; + private final InMemoryArtifactService delegate = new InMemoryArtifactService(); + private final List saves = new ArrayList<>(); + + private FakeArtifactService( + String uriPrefix, + ImmutableSet failingNames, + ImmutableSet erroringNames, + boolean versionsReadable) { + this.uriPrefix = uriPrefix; + this.failingNames = failingNames; + this.erroringNames = erroringNames; + this.versionsReadable = versionsReadable; + } + + static FakeArtifactService create() { + return new FakeArtifactService(GS_PREFIX, ImmutableSet.of(), ImmutableSet.of(), true); + } + + static FakeArtifactService withUriPrefix(String uriPrefix) { + return new FakeArtifactService(uriPrefix, ImmutableSet.of(), ImmutableSet.of(), true); + } + + static FakeArtifactService failingOn(String fileName) { + return new FakeArtifactService(GS_PREFIX, ImmutableSet.of(fileName), ImmutableSet.of(), true); + } + + static FakeArtifactService erroringOn(String fileName) { + return new FakeArtifactService(GS_PREFIX, ImmutableSet.of(), ImmutableSet.of(fileName), true); + } + + static FakeArtifactService withUnreadableVersions() { + return new FakeArtifactService(GS_PREFIX, ImmutableSet.of(), ImmutableSet.of(), false); + } + + ImmutableList savedNames() { + return saves.stream().map(Save::fileName).collect(toImmutableList()); + } + + @Override + public Single saveArtifact( + String appName, String userId, String sessionId, String filename, Part artifact) { + if (failingNames.contains(filename)) { + throw new IllegalStateException("Storage error"); + } + if (erroringNames.contains(filename)) { + return Single.error(() -> new IllegalStateException("Storage error")); + } + saves.add(new Save(filename, artifact)); + return delegate.saveArtifact(appName, userId, sessionId, filename, artifact); + } + + @Override + public Single> listVersions( + String appName, String userId, String sessionId, String filename) { + if (!versionsReadable) { + throw new IllegalStateException("Storage error"); + } + return delegate.listVersions(appName, userId, sessionId, filename); + } + + @Override + public Single saveAndReloadArtifact( + String appName, String userId, String sessionId, String filename, Part artifact) { + return saveArtifact(appName, userId, sessionId, filename, artifact) + .map(version -> storedAt(STORED_URI.formatted(uriPrefix, filename, version))); + } + + @Override + public Maybe loadArtifact( + String appName, + String userId, + String sessionId, + String filename, + @Nullable Integer version) { + return delegate.loadArtifact(appName, userId, sessionId, filename, version); + } + + @Override + public Single listArtifactKeys( + String appName, String userId, String sessionId) { + return delegate.listArtifactKeys(appName, userId, sessionId); + } + + @Override + public Completable deleteArtifact( + String appName, String userId, String sessionId, String filename) { + return delegate.deleteArtifact(appName, userId, sessionId, filename); + } + + private static Part storedAt(String uri) { + return Part.builder() + .fileData(FileData.builder().mimeType("text/plain").fileUri(uri)) + .build(); + } + } + + /** Saves nothing: returns version 0, reloads the artifact as given, and lists no versions. */ + private static final class DummyArtifactService implements BaseArtifactService { + @Override + public Single saveArtifact( + String appName, String userId, String sessionId, String filename, Part artifact) { + return Single.just(0); + } + + @Override + public Single saveAndReloadArtifact( + String appName, String userId, String sessionId, String filename, Part artifact) { + return Single.just(artifact); + } + + @Override + public Maybe loadArtifact( + String appName, + String userId, + String sessionId, + String filename, + @Nullable Integer version) { + return Maybe.empty(); + } + + @Override + public Single listArtifactKeys( + String appName, String userId, String sessionId) { + return Single.just(ListArtifactsResponse.builder().filenames(ImmutableList.of()).build()); + } + + @Override + public Completable deleteArtifact( + String appName, String userId, String sessionId, String filename) { + return Completable.complete(); + } + + @Override + public Single> listVersions( + String appName, String userId, String sessionId, String filename) { + return Single.just(ImmutableList.of()); + } + } +}