-
Notifications
You must be signed in to change notification settings - Fork 405
fix(artifacts): honor the "user:" namespace in InMemoryArtifactService #1449
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
| * | ||
|
|
@@ -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); | ||
|
|
@@ -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) | ||
| .computeIfAbsent(filename, unused -> new ArrayList<>()); | ||
|
|
||
| if (versions.isEmpty()) { | ||
|
|
@@ -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()); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Python ends How about |
||
| } | ||
|
|
||
| /** | ||
|
|
@@ -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(); | ||
| } | ||
|
|
||
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I checked and this one leaks across sessions as well. Python's |
||
| if (size == 0) { | ||
|
|
@@ -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 |
|---|---|---|
|
|
@@ -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; | ||
|
|
@@ -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; | ||
|
|
||
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 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(); | ||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In a read path
computeIfAbsentwrites an empty entry on a miss, andlistArtifactKeysreportskeySet(). 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") -> emptylistArtifactKeys(..., 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