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
Expand Up @@ -26,19 +26,36 @@
import io.reactivex.rxjava3.core.Single;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.IntStream;
import org.jspecify.annotations.Nullable;

/** An in-memory implementation of the {@link BaseArtifactService}. */
public final class InMemoryArtifactService implements BaseArtifactService {
// Mirrors the "user" path segment GcsArtifactService.getBlobPrefix uses for user-namespaced
// filenames: a stand-in session key so those artifacts are stored and looked up independently
// of whatever sessionId happened to be current when they were saved.
private static final String USER_NAMESPACE_SESSION_KEY = "user";

private final Map<String, Map<String, Map<String, Map<String, List<Part>>>>> artifacts;

public InMemoryArtifactService() {
this.artifacts = new HashMap<>();
}

/**
* Checks if a filename uses the user namespace.
*
* @param filename Filename to check.
* @return true if prefixed with "user:", false otherwise.
*/
private boolean fileHasUserNamespace(String filename) {
return filename != null && filename.startsWith("user:");
}

/**
* Saves an artifact in memory and assigns a new version.
*
Expand All @@ -48,7 +65,7 @@ public InMemoryArtifactService() {
public Single<Integer> saveArtifact(
String appName, String userId, String sessionId, String filename, Part artifact) {
List<Part> versions =
getArtifactsMap(appName, userId, sessionId)
getArtifactsMap(appName, userId, sessionId, filename)
.computeIfAbsent(filename, unused -> new ArrayList<>());
versions.add(artifact);
return Single.just(versions.size() - 1);
Expand All @@ -63,7 +80,7 @@ public Single<Integer> saveArtifact(
public Maybe<Part> loadArtifact(
String appName, String userId, String sessionId, String filename, @Nullable Integer version) {
List<Part> versions =
getArtifactsMap(appName, userId, sessionId)
getArtifactsMap(appName, userId, sessionId, filename)

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.

In a read path computeIfAbsent writes an empty entry on a miss, and listArtifactKeys reports keySet(). So a failed load invents a file. For reference, Python's read is pure (self.artifacts.get(path)).

This is a pre-existing issue — except this PR changes its reach. On main the ghost stays in the caller's session; now user: files share one bucket, so every session of that user sees it:

loadArtifact(..., sessionA, "user:missing.json") -> empty
listArtifactKeys(..., sessionB) -> ["user:missing.json"] (main: [])

What do you think of going back to a pure read, ImmutableList.of() is a shared singleton, so it allocates nothing

List<Part> versions = 
  getArtifactsMap(appName, userId, sessionId, filename)
    .getOrDefault(filename, ImmutableList.of());

.computeIfAbsent(filename, unused -> new ArrayList<>());

if (versions.isEmpty()) {
Expand All @@ -88,10 +105,15 @@ public Maybe<Part> loadArtifact(
@Override
public Single<ListArtifactsResponse> listArtifactKeys(
String appName, String userId, String sessionId) {
Set<String> filenames = new HashSet<>();
// Session-scoped filenames, keyed by the actual sessionId.
filenames.addAll(getArtifactsMapForSessionKey(appName, userId, sessionId).keySet());
// User-namespaced filenames live under the shared USER_NAMESPACE_SESSION_KEY bucket
// regardless of sessionId, mirroring GcsArtifactService's merged session+user listing.
filenames.addAll(
getArtifactsMapForSessionKey(appName, userId, USER_NAMESPACE_SESSION_KEY).keySet());
return Single.just(
ListArtifactsResponse.builder()
.filenames(ImmutableList.copyOf(getArtifactsMap(appName, userId, sessionId).keySet()))
.build());
ListArtifactsResponse.builder().filenames(ImmutableList.copyOf(filenames)).build());

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.

Python ends list_artifact_keys with return sorted(filenames). ImmutableList.copyOf over a HashSet gives arbitrary order instead. Additionally, Java's own GcsArtifactService already sorts (:186), so the two backends currently disagree.

How about ListArtifactsResponse.builder().filenames(ImmutableList.sortedCopyOf(filenames)).build()); ?

}

/**
Expand All @@ -102,7 +124,7 @@ public Single<ListArtifactsResponse> listArtifactKeys(
@Override
public Completable deleteArtifact(
String appName, String userId, String sessionId, String filename) {
getArtifactsMap(appName, userId, sessionId).remove(filename);
getArtifactsMap(appName, userId, sessionId, filename).remove(filename);
return Completable.complete();
}

Expand All @@ -115,7 +137,7 @@ public Completable deleteArtifact(
public Single<ImmutableList<Integer>> listVersions(
String appName, String userId, String sessionId, String filename) {
int size =
getArtifactsMap(appName, userId, sessionId)
getArtifactsMap(appName, userId, sessionId, filename)
.computeIfAbsent(filename, unused -> new ArrayList<>())
.size();
Comment on lines 139 to 142

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.

I checked and this one leaks across sessions as well. Python's list_versions is also a pure read. Same fix could apply:

int size =
  getArtifactsMap(appName, userId, sessionId, filename)
    .getOrDefault(filename, ImmutableList.of())
    .size();

if (size == 0) {
Expand All @@ -131,10 +153,22 @@ public Single<Part> saveAndReloadArtifact(
.flatMap(version -> loadArtifact(appName, userId, sessionId, filename, version).toSingle());
}

private Map<String, List<Part>> getArtifactsMap(String appName, String userId, String sessionId) {
/**
* Resolves the artifacts map for a filename, routing user-namespaced filenames to the shared
* {@link #USER_NAMESPACE_SESSION_KEY} bucket instead of the given session, so they are visible
* across all of a user's sessions exactly as {@link GcsArtifactService} stores them.
*/
private Map<String, List<Part>> getArtifactsMap(
String appName, String userId, String sessionId, String filename) {
String sessionKey = fileHasUserNamespace(filename) ? USER_NAMESPACE_SESSION_KEY : sessionId;
return getArtifactsMapForSessionKey(appName, userId, sessionKey);
}

private Map<String, List<Part>> getArtifactsMapForSessionKey(
String appName, String userId, String sessionKey) {
return artifacts
.computeIfAbsent(appName, unused -> new HashMap<>())
.computeIfAbsent(userId, unused -> new HashMap<>())
.computeIfAbsent(sessionId, unused -> new HashMap<>());
.computeIfAbsent(sessionKey, unused -> new HashMap<>());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

import static com.google.common.truth.Truth.assertThat;

import com.google.common.collect.ImmutableList;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
Expand All @@ -34,6 +35,9 @@ public class InMemoryArtifactServiceTest {
private static final String USER_ID = "test-user";
private static final String SESSION_ID = "test-session";
private static final String FILENAME = "test-file.txt";
private static final String USER_FILENAME = "user:config.json";
private static final String SESSION_A = "session-A";
private static final String SESSION_B = "session-B";

private InMemoryArtifactService service;

Expand Down Expand Up @@ -72,6 +76,84 @@ public void saveAndReloadArtifact_reloadsArtifact() {
assertThat(result).hasValue(artifact);
}

@Test
public void save_userNamespace_visibleAcrossSessions() {
// "user:"-prefixed filenames are documented/tested (GcsArtifactServiceTest) as living in a
// session-independent namespace: GcsArtifactService.getBlobPrefix special-cases
// fileHasUserNamespace(filename) to build "appName/userId/user/filename/" (no sessionId
// segment at all). InMemoryArtifactService must mirror that: a "user:" artifact saved in one
// session has to be readable from a different session for the same app/user.
Part artifact = Part.fromBytes(new byte[] {9, 9, 9}, "application/json");

var unused =
service.saveArtifact(APP_NAME, USER_ID, SESSION_A, USER_FILENAME, artifact).blockingGet();

Optional<Part> resultFromOtherSession =
asOptional(service.loadArtifact(APP_NAME, USER_ID, SESSION_B, USER_FILENAME));

assertThat(resultFromOtherSession).hasValue(artifact);
}

@Test
public void load_nonNamespacedFilename_notVisibleAcrossSessions() {
// Regression guard: the user-namespace fix must not make every artifact global. A
// non-prefixed filename saved in one session must remain invisible from a different session,
// exactly as before the fix.
Part artifact = Part.fromBytes(new byte[] {1, 2, 3}, "text/plain");

var unused =
service.saveArtifact(APP_NAME, USER_ID, SESSION_A, FILENAME, artifact).blockingGet();

Optional<Part> resultFromOtherSession =
asOptional(service.loadArtifact(APP_NAME, USER_ID, SESSION_B, FILENAME));

assertThat(resultFromOtherSession).isEmpty();
}

@Test
public void listArtifactKeys_userNamespace_visibleAcrossSessions() {
Part artifact = Part.fromBytes(new byte[] {9, 9, 9}, "application/json");

var unused =
service.saveArtifact(APP_NAME, USER_ID, SESSION_A, USER_FILENAME, artifact).blockingGet();

ListArtifactsResponse response =
service.listArtifactKeys(APP_NAME, USER_ID, SESSION_B).blockingGet();

assertThat(response.filenames()).contains(USER_FILENAME);
}

Comment on lines +113 to +125

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.

This test passes when the list would be full of ghosts so it might be worth adding one more, proving a failed load leaves nothing behind

@Test
public void listArtifactKeys_failedUserNamespaceLoad_doesNotCreatePhantomKey() {

Optional<Part> missing = 
  asOptional(service.loadArtifact(APP_NAME, USER_ID, SESSION_A, USER_FILENAME));

assertThat(missing).isEmpty();

ListArtifactsResponse response = 
  service.listArtifactKeys(APP_NAME, USER_ID, SESSION_B).blockingGet();

assertThat(response.filenames()).isEmpty();

}

and possibly similarly for listVersions?

@Test
public void listVersions_userNamespace_visibleAcrossSessions() {
Part artifact1 = Part.fromBytes(new byte[] {1}, "application/json");
Part artifact2 = Part.fromBytes(new byte[] {1, 2}, "application/json");

var unused1 =
service.saveArtifact(APP_NAME, USER_ID, SESSION_A, USER_FILENAME, artifact1).blockingGet();
var unused2 =
service.saveArtifact(APP_NAME, USER_ID, SESSION_A, USER_FILENAME, artifact2).blockingGet();

ImmutableList<Integer> versions =
service.listVersions(APP_NAME, USER_ID, SESSION_B, USER_FILENAME).blockingGet();

assertThat(versions).containsExactly(0, 1).inOrder();
}

@Test
public void deleteArtifact_userNamespace_removesAcrossSessions() {
Part artifact = Part.fromBytes(new byte[] {9, 9, 9}, "application/json");

var unused =
service.saveArtifact(APP_NAME, USER_ID, SESSION_A, USER_FILENAME, artifact).blockingGet();

service.deleteArtifact(APP_NAME, USER_ID, SESSION_B, USER_FILENAME).blockingAwait();

Optional<Part> resultFromOriginalSession =
asOptional(service.loadArtifact(APP_NAME, USER_ID, SESSION_A, USER_FILENAME));

assertThat(resultFromOriginalSession).isEmpty();
}

private static <T> Optional<T> asOptional(Maybe<T> maybe) {
return maybe.map(Optional::of).defaultIfEmpty(Optional.empty()).blockingGet();
}
Expand Down