This is an automated email from the ASF dual-hosted git repository. voidmatcha pushed a commit to branch temp/seung-00/assistant-transport-base in repository https://gitbox.apache.org/repos/asf/zeppelin.git
commit 6ce4f4bfdf1b122cefb3a574611e0053168795fc Author: YONGJAE LEE <[email protected]> AuthorDate: Thu Oct 8 21:10:07 2026 +0900 Validate assistant conversation storage --- .../service/assistant/ConversationJsonCodec.java | 41 +++++++++- .../assistant/FileConversationRepository.java | 12 ++- .../assistant/FileConversationRepositoryTest.java | 88 ++++++++++++++++++++++ 3 files changed, 137 insertions(+), 4 deletions(-) diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationJsonCodec.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationJsonCodec.java index c4d3776f1a..8ec38567d4 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationJsonCodec.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationJsonCodec.java @@ -28,6 +28,8 @@ import com.google.gson.JsonSerializationContext; import com.google.gson.JsonSerializer; import java.lang.reflect.Type; +import java.time.Instant; +import java.time.format.DateTimeParseException; /** * JSON wire format for an assistant conversation. @@ -48,9 +50,20 @@ final class ConversationJsonCodec { static Conversation deserialize(String json) { Conversation conversation = GSON.fromJson(json, Conversation.class); if (conversation == null || conversation.getId() == null || conversation.getNoteId() == null + || conversation.getOwnerId() == null || conversation.getOwnerId().isBlank() || conversation.getCreatedAt() == null || conversation.getMessages() == null) { throw new JsonParseException("Invalid assistant conversation storage"); } + try { + Instant.parse(conversation.getCreatedAt()); + } catch (DateTimeParseException e) { + throw new JsonParseException("Invalid assistant conversation creation time"); + } + for (Message message : conversation.getMessages()) { + if (message == null) { + throw new JsonParseException("Invalid null assistant message storage"); + } + } return conversation; } @@ -64,23 +77,45 @@ final class ConversationJsonCodec { @Override public Message deserialize(JsonElement json, Type type, JsonDeserializationContext ctx) throws JsonParseException { + if (!json.isJsonObject()) { + throw new JsonParseException("Invalid assistant message storage"); + } JsonObject obj = json.getAsJsonObject(); - if (!obj.has("role")) { - throw new JsonParseException("Message missing 'role' field"); + requireString(obj, "id", false); + requireString(obj, "content", true); + JsonElement roleElement = obj.get("role"); + if (roleElement == null || !roleElement.isJsonPrimitive() + || !roleElement.getAsJsonPrimitive().isString()) { + throw new JsonParseException("Message missing or invalid 'role' field"); + } + Message.Role role; + try { + role = Message.Role.fromValue(roleElement.getAsString()); + } catch (IllegalArgumentException e) { + throw new JsonParseException("Invalid assistant message role"); } - Message.Role role = Message.Role.fromValue(obj.get("role").getAsString()); switch (role) { case USER: return SUBTYPE.fromJson(json, Message.User.class); case ASSISTANT: return SUBTYPE.fromJson(json, Message.Assistant.class); case TOOL: + requireString(obj, "toolCallId", false); return SUBTYPE.fromJson(json, Message.Tool.class); default: throw new JsonParseException("Unsupported role: " + role); } } + private static void requireString(JsonObject obj, String field, boolean allowEmpty) { + JsonElement value = obj.get(field); + if (value == null || !value.isJsonPrimitive() + || !value.getAsJsonPrimitive().isString() + || (!allowEmpty && value.getAsString().isBlank())) { + throw new JsonParseException("Message missing or invalid '" + field + "' field"); + } + } + @Override public JsonElement serialize(Message src, Type type, JsonSerializationContext ctx) { return SUBTYPE.toJsonTree(src); diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/FileConversationRepository.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/FileConversationRepository.java index b9f371b33e..ca6b18da8f 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/FileConversationRepository.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/FileConversationRepository.java @@ -17,6 +17,7 @@ package org.apache.zeppelin.service.assistant; +import com.google.gson.JsonParseException; import java.io.File; import java.io.IOException; import java.nio.file.Files; @@ -29,8 +30,11 @@ import java.util.List; import java.util.Optional; import java.util.function.Consumer; import org.apache.zeppelin.util.FileUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public final class FileConversationRepository implements ConversationRepository { + private static final Logger LOGGER = LoggerFactory.getLogger(FileConversationRepository.class); private final File directory; public FileConversationRepository(File directory) { @@ -45,7 +49,13 @@ public final class FileConversationRepository implements ConversationRepository if (files == null) throw new IOException(); List<Conversation> conversations = new ArrayList<>(); - for (File file : files) read(file).ifPresent(conversations::add); + for (File file : files) { + try { + read(file).ifPresent(conversations::add); + } catch (JsonParseException e) { + LOGGER.warn("Skipping invalid assistant conversation file {}", file.getName()); + } + } conversations.sort(Comparator.comparing(c -> Instant.parse(c.getCreatedAt()))); return Collections.unmodifiableList(conversations); } diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/FileConversationRepositoryTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/FileConversationRepositoryTest.java index efeaa94b32..eae7ba251b 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/FileConversationRepositoryTest.java +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/FileConversationRepositoryTest.java @@ -19,9 +19,12 @@ package org.apache.zeppelin.service.assistant; import static org.junit.jupiter.api.Assertions.*; +import com.google.gson.JsonParseException; +import com.google.gson.JsonParser; import java.io.File; import java.io.IOException; import java.nio.file.Files; +import java.util.List; import java.util.Set; import java.util.stream.Collectors; import org.junit.jupiter.api.Test; @@ -52,6 +55,91 @@ class FileConversationRepositoryTest { var sut = new FileConversationRepository(directory); Files.createDirectories(new File(directory, "noteId/unreadable.json").toPath()); assertThrows(IOException.class, () -> sut.findAll("noteId")); + assertThrows(IOException.class, () -> sut.find("noteId", "unreadable")); + } + + @Test + void skipsMalformedSiblingWithoutTreatingItAsMissing(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var conversation = Conversation.create("noteId", "valid", "holden"); + sut.create(conversation); + Files.writeString(new File(directory, "noteId/corrupt.json").toPath(), "{incomplete"); + + assertEquals(List.of(conversation.getId()), sut.findAll("noteId").stream() + .map(Conversation::getId).collect(Collectors.toList())); + assertThrows(JsonParseException.class, () -> sut.find("noteId", "corrupt")); + } + + @Test + void rejectsMissingOwnerAndInvalidCreationTime(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var conversation = Conversation.create("noteId", "valid", "holden"); + sut.create(conversation); + + for (String field : List.of("ownerId", "createdAt")) { + var json = JsonParser.parseString(ConversationJsonCodec.serialize(conversation)) + .getAsJsonObject(); + json.remove(field); + Files.writeString(new File(directory, "noteId/missing-" + field + ".json").toPath(), + json.toString()); + assertThrows(JsonParseException.class, () -> sut.find("noteId", "missing-" + field)); + + json.addProperty(field, ""); + Files.writeString(new File(directory, "noteId/invalid-" + field + ".json").toPath(), + json.toString()); + assertThrows(JsonParseException.class, () -> sut.find("noteId", "invalid-" + field)); + } + + assertEquals(List.of(conversation.getId()), sut.findAll("noteId").stream() + .map(Conversation::getId).collect(Collectors.toList())); + } + + @Test + void skipsSiblingWithInvalidMessageRole(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var conversation = Conversation.create("noteId", "valid", "holden"); + sut.create(conversation); + conversation.addMessage(Message.user("msg", "content")); + var json = JsonParser.parseString(ConversationJsonCodec.serialize(conversation)) + .getAsJsonObject(); + json.getAsJsonArray("messages").get(0).getAsJsonObject().addProperty("role", "unknown"); + Files.writeString(new File(directory, "noteId/invalid-role.json").toPath(), json.toString()); + + assertEquals(List.of(conversation.getId()), sut.findAll("noteId").stream() + .map(Conversation::getId).collect(Collectors.toList())); + assertThrows(JsonParseException.class, () -> sut.find("noteId", "invalid-role")); + } + + @Test + void rejectsNullAndIncompleteMessages(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var valid = Conversation.create("noteId", "valid", "holden"); + valid.addMessage(Message.user("user", "hello")); + valid.addMessage(Message.assistant("assistant", "")); + valid.addMessage(Message.tool("tool", "call", "")); + sut.create(valid); + + var invalidMessages = List.of("null", "{}", + "{\"role\":\"user\",\"content\":\"hello\"}", + "{\"role\":\"user\",\"id\":\" \",\"content\":\"hello\"}", + "{\"role\":\"user\",\"id\":1,\"content\":\"hello\"}", + "{\"role\":\"user\",\"id\":\"user\"}", + "{\"role\":\"assistant\",\"id\":\"assistant\",\"content\":null}", + "{\"role\":\"assistant\",\"id\":\"assistant\",\"content\":1}", + "{\"role\":\"tool\",\"id\":\"tool\",\"content\":\"result\"}"); + for (int i = 0; i < invalidMessages.size(); i++) { + var json = JsonParser.parseString(ConversationJsonCodec.serialize(valid)).getAsJsonObject(); + json.getAsJsonArray("messages").set(0, JsonParser.parseString(invalidMessages.get(i))); + String id = "invalid-message-" + i; + Files.writeString(new File(directory, "noteId/" + id + ".json").toPath(), json.toString()); + assertThrows(JsonParseException.class, () -> sut.find("noteId", id), id); + } + + assertEquals(List.of(valid.getId()), sut.findAll("noteId").stream() + .map(Conversation::getId).collect(Collectors.toList())); + var loaded = sut.find("noteId", valid.getId()).orElseThrow(); + assertEquals("", ((Message.Assistant) loaded.getMessages().get(1)).getContent()); + assertEquals("", ((Message.Tool) loaded.getMessages().get(2)).getContent()); } @Test
