Validate and parameterize Neo4j chat memory labels

memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java

@@ -46,6 +46,7 @@
  *
  * @author Enrico Rampazzo
  * @author Michael J. Simons
+ * @author Soby Chacko
  * @since 1.0.0
  */
 
@@ -138,30 +139,29 @@ public void saveAll(String conversationId, List<Message> messages) {
 
 	@Override
 	public void deleteByConversationId(String conversationId) {
-		// First delete all messages and related nodes
 		String deleteMessagesStatement = """
-				MATCH (s:%s {id:$conversationId})-[r:HAS_MESSAGE]->(m:%s)
-				OPTIONAL MATCH (m)-[:HAS_METADATA]->(metadata:%s)
-				OPTIONAL MATCH (m)-[:HAS_MEDIA]->(media:%s)
-				OPTIONAL MATCH (m)-[:HAS_TOOL_RESPONSE]-(tr:%s)
-				OPTIONAL MATCH (m)-[:HAS_TOOL_CALL]->(tc:%s)
+				MATCH (s:$($sessionLabel) {id:$conversationId})-[r:HAS_MESSAGE]->(m:$($messageLabel))
+				OPTIONAL MATCH (m)-[:HAS_METADATA]->(metadata:$($metadataLabel))
+				OPTIONAL MATCH (m)-[:HAS_MEDIA]->(media:$($mediaLabel))
+				OPTIONAL MATCH (m)-[:HAS_TOOL_RESPONSE]-(tr:$($toolResponseLabel))
+				OPTIONAL MATCH (m)-[:HAS_TOOL_CALL]->(tc:$($toolCallLabel))
 				DETACH DELETE m, metadata, media, tr, tc
-				""".formatted(this.config.getSessionLabel(), this.config.getMessageLabel(),
-				this.config.getMetadataLabel(), this.config.getMediaLabel(), this.config.getToolResponseLabel(),
-				this.config.getToolCallLabel());
+				""";
 
-		// Then delete the conversation node itself
 		String deleteConversationStatement = """
-				MATCH (s:%s {id:$conversationId})
+				MATCH (s:$($sessionLabel) {id:$conversationId})
 				DETACH DELETE s
-				""".formatted(this.config.getSessionLabel());
+				""";
+
+		Map<String, Object> params = Map.of("conversationId", conversationId, "sessionLabel",
+				this.config.getSessionLabel(), "messageLabel", this.config.getMessageLabel(), "metadataLabel",
+				this.config.getMetadataLabel(), "mediaLabel", this.config.getMediaLabel(), "toolResponseLabel",
+				this.config.getToolResponseLabel(), "toolCallLabel", this.config.getToolCallLabel());
 
 		try (Session s = this.config.getDriver().session()) {
 			try (Transaction t = s.beginTransaction()) {
-				// First delete messages
-				t.run(deleteMessagesStatement, Map.of("conversationId", conversationId));
-				// Then delete the conversation node
-				t.run(deleteConversationStatement, Map.of("conversationId", conversationId));
+				t.run(deleteMessagesStatement, params);
+				t.run(deleteConversationStatement, params);
 				t.commit();
 			}
 		}

memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryConfig.java

@@ -16,6 +16,8 @@
 
 package org.springframework.ai.chat.memory.repository.neo4j;
 
+import java.util.regex.Pattern;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.jspecify.annotations.Nullable;
@@ -27,6 +29,7 @@
  * Configuration for the Neo4j Chat Memory store.
  *
  * @author Enrico Rampazzo
+ * @author Soby Chacko
  */
 public final class Neo4jChatMemoryRepositoryConfig {
 
@@ -46,6 +49,8 @@ public final class Neo4jChatMemoryRepositoryConfig {
 
 	private static final Log logger = LogFactory.getLog(Neo4jChatMemoryRepositoryConfig.class);
 
+	private static final Pattern SAFE_LABEL = Pattern.compile("[\\p{Alpha}_][\\p{Alnum}_]*");
+
 	private final Driver driver;
 
 	private final String sessionLabel;
@@ -97,9 +102,20 @@ private Neo4jChatMemoryRepositoryConfig(Builder builder) {
 		this.toolCallLabel = builder.toolCallLabel;
 		this.metadataLabel = builder.metadataLabel;
 		this.toolResponseLabel = builder.toolResponseLabel;
+		validateLabels();
 		ensureIndexes();
 	}
 
+	private void validateLabels() {
+		for (String label : new String[] { this.sessionLabel, this.messageLabel, this.metadataLabel, this.mediaLabel,
+				this.toolCallLabel, this.toolResponseLabel }) {
+			if (!SAFE_LABEL.matcher(label).matches()) {
+				throw new IllegalArgumentException("Invalid Neo4j node label: '" + label
+						+ "'. Labels must start with a letter or underscore and contain only letters, digits, or underscores.");
+			}
+		}
+	}
+
 	/**
 	 * Ensures that indexes exist on conversationId for Session nodes and index for
 	 * Message nodes. This improves query performance for lookups and ordering.

memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4JChatMemoryRepositoryConfigIT.java

@@ -16,9 +16,15 @@
 
 package org.springframework.ai.chat.memory.repository.neo4j;
 
+import java.util.function.BiFunction;
+import java.util.stream.Stream;
+
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.GraphDatabase;
 import org.neo4j.driver.Result;
@@ -28,7 +34,13 @@
 import org.testcontainers.junit.jupiter.Testcontainers;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
+/**
+ * Integration tests for {@link Neo4jChatMemoryRepositoryConfig}.
+ *
+ * @author Soby Chacko
+ */
 @Testcontainers
 class Neo4JChatMemoryRepositoryConfigIT {
 
@@ -87,6 +99,34 @@ void builderShouldSetCustomLabels() {
 		assertThat(config.getMessageLabel()).isEqualTo(customMessageLabel);
 	}
 
+	@ParameterizedTest
+	@ValueSource(strings = { "1Session", "Session Label", "Session-Label", "Session`", "Session)" })
+	void shouldRejectMalformedSessionLabel(String invalidLabel) {
+		assertThatThrownBy(() -> Neo4jChatMemoryRepositoryConfig.builder()
+			.withDriver(driver)
+			.withSessionLabel(invalidLabel)
+			.build()).isInstanceOf(IllegalArgumentException.class).hasMessageContaining("Invalid Neo4j node label");
+	}
+
+	static Stream<BiFunction<Neo4jChatMemoryRepositoryConfig.Builder, String, Neo4jChatMemoryRepositoryConfig.Builder>> allLabelSetters() {
+		return Stream.of(Neo4jChatMemoryRepositoryConfig.Builder::withSessionLabel,
+				Neo4jChatMemoryRepositoryConfig.Builder::withMessageLabel,
+				Neo4jChatMemoryRepositoryConfig.Builder::withMetadataLabel,
+				Neo4jChatMemoryRepositoryConfig.Builder::withMediaLabel,
+				Neo4jChatMemoryRepositoryConfig.Builder::withToolCallLabel,
+				Neo4jChatMemoryRepositoryConfig.Builder::withToolResponseLabel);
+	}
+
+	@ParameterizedTest
+	@MethodSource("allLabelSetters")
+	void shouldRejectSpecialCharactersInAnyLabel(
+			BiFunction<Neo4jChatMemoryRepositoryConfig.Builder, String, Neo4jChatMemoryRepositoryConfig.Builder> setter) {
+		assertThatThrownBy(
+				() -> setter.apply(Neo4jChatMemoryRepositoryConfig.builder().withDriver(driver), "Bad`) Label").build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("Invalid Neo4j node label");
+	}
+
 	@Test
 	void gettersShouldReturnConfiguredValues() {
 		Neo4jChatMemoryRepositoryConfig config = Neo4jChatMemoryRepositoryConfig.builder()

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat-memory.adoc

@@ -364,6 +364,9 @@ ChatMemory chatMemory = MessageWindowChatMemory.builder()
 | `spring.ai.chat.memory.repository.neo4j.media-label` | The label for the nodes that store media associated with a message | `Media`
 |===
 
+NOTE: Label values must be valid Neo4j identifiers: they must start with a letter or underscore and contain only letters, digits, or underscores.
+Supplying a value that does not meet these requirements will cause an `IllegalArgumentException` to be thrown at application startup.
+
 ==== Index Initialization
 
 The Neo4j repository will automatically ensure that indexes are created for conversation IDs and message indices to optimize performance. If you use custom labels, indexes will be created for those labels as well. No schema initialization is required, but you should ensure your Neo4j instance is accessible to your application.