feat: Enhance JsonAgentResponseHandler with JSON sanitization and add comprehensive unit tests

src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java

@@ -66,10 +66,19 @@ public class JsonAgentResponseHandler implements AgentResponseHandler {
 			jsonNode = this.objectMapper.readTree(response);
 		}
 		catch (final Exception e) {
-			LOGGER.warn("Cannot parse json response: {}", completion.response());
+			LOGGER.debug("Initial JSON parse failed, attempting to sanitize: {}", e.getMessage());
+			try {
+				final var response = completion.response().substring(firstCurly, lastCurly + 1);
+				final var sanitized = sanitizeJson(response);
+				jsonNode = this.objectMapper.readTree(sanitized);
+				LOGGER.debug("Successfully parsed after sanitization");
+			}
+			catch (final Exception e2) {
+				LOGGER.warn("Cannot parse json response even after sanitization: {}", completion.response());
 				addJsonResponseErrorPrompt(context);
 				return false;
 			}
+		}
 		
 		final var items = new LinkedHashMap<String, List<String>>();
 		final Iterator<Map.Entry<String, JsonNode>> fields = jsonNode.properties().iterator();
@@ -103,12 +112,8 @@ public class JsonAgentResponseHandler implements AgentResponseHandler {
 	
 	private void parseProperty(final LinkedHashMap<String, List<String>> items, final Entry<String, JsonNode> nestedProperty) {
 		if (nestedProperty.getValue().isObject()) {
-			final var nestedValuesIter = nestedProperty.getValue().values();
+			for (final var entry : nestedProperty.getValue().properties()) {
-			while (nestedValuesIter.hasNext()) {
+				this.parseProperty(items, entry);
-				final var v = nestedValuesIter.next();
-				if (v.isTextual() || v.isNumber() || v.isBoolean()) {
-					items.computeIfAbsent(this.prefixFields + nestedProperty.getKey(), k -> new ArrayList<>()).add(v.asText());
-				}
 			}
 		}
 		else {
@@ -121,4 +126,61 @@ public class JsonAgentResponseHandler implements AgentResponseHandler {
 		context.getConversation().replaceTurn(new Conversation.Turn(turn.type(), turn.text() + "\n\nThe previous response could not be fully processed or validated. Please make sure to format the response in valid JSON syntax with properly escaped characters."));
 	}
 	
+	private String sanitizeJson(final String json) {
+		final var result = new StringBuilder(json.length() + 10);
+		boolean inString = false;
+		boolean escaped = false;
+		
+		for (int i = 0; i < json.length(); i++) {
+			final char c = json.charAt(i);
+			
+			if (escaped) {
+				result.append(c);
+				escaped = false;
+				continue;
+			}
+			
+			if (c == '\\' && inString) {
+				escaped = true;
+				result.append(c);
+				continue;
+			}
+			
+			if (c == '"') {
+				boolean unescapedQuote = false;
+				if (inString) {
+					for (int j = i + 1; j < json.length(); j++) {
+						final var nextChar = json.charAt(j);
+						if (Character.isWhitespace(nextChar)) {
+							continue;
+						}
+						if (nextChar == ':' || nextChar == ',' || nextChar == '}' || nextChar == ']') {
+							break;
+						}
+						result.append('\\');
+						unescapedQuote = true;
+						break;
+					}
+				}
+				
+				inString = unescapedQuote || !inString;
+				result.append(c);
+				continue;
+			}
+			
+			if (inString && (c == '}' || c == ']') && (i == 0 || json.charAt(i - 1) != '\\')) {
+				int insertPos = result.length();
+				while (insertPos > 0 && Character.isWhitespace(result.charAt(insertPos - 1))) {
+					insertPos--;
+				}
+				result.insert(insertPos, '"');
+				inString = false;
+			}
+			
+			result.append(c);
+		}
+		
+		return result.toString();
+	}
+	
 }

src/test/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandlerTest.java

@@ -0,0 +1,369 @@
+package de.extio.lmlib.agent.responsehandler;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+
+import de.extio.lmlib.agent.AgentContext;
+import de.extio.lmlib.client.Completion;
+import de.extio.lmlib.client.Conversation;
+
+public class JsonAgentResponseHandlerTest {
+	
+	@Test
+	public void testSimpleJsonParsing() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"name": "John Doe",
+					"age": "30",
+					"city": "New York"
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertEquals("John Doe", context.getStringValue("name"));
+		assertEquals("30", context.getStringValue("age"));
+		assertEquals("New York", context.getStringValue("city"));
+	}
+	
+	@Test
+	public void testJsonWithPrefix() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler("test_");
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"field1": "value1",
+					"field2": "value2"
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertEquals("value1", context.getStringValue("test_field1"));
+		assertEquals("value2", context.getStringValue("test_field2"));
+	}
+	
+	@Test
+	public void testJsonWithPreamble() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				Here is the result:
+				
+				{
+					"status": "success",
+					"message": "Operation completed"
+				}
+				
+				Additional notes after the JSON.
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertEquals("success", context.getStringValue("status"));
+		assertEquals("Operation completed", context.getStringValue("message"));
+	}
+	
+	@Test
+	public void testJsonArray() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"items": ["apple", "banana", "cherry"]
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		final List<String> items = context.getStringValues("items");
+		assertNotNull(items);
+		assertEquals(3, items.size());
+		assertEquals("apple", items.get(0));
+		assertEquals("banana", items.get(1));
+		assertEquals("cherry", items.get(2));
+	}
+	
+	@Test
+	public void testJsonNestedObject() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"user": {
+						"name": "Jane",
+						"email": "jane@example.com"
+					}
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertEquals("Jane", context.getStringValue("name"));
+		assertEquals("jane@example.com", context.getStringValue("email"));
+	}
+	
+	@Test
+	public void testJsonArrayOfObjects() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"users": [
+						{"name": "Alice", "age": "25"},
+						{"name": "Bob", "age": "30"}
+					]
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		final List<String> names = context.getStringValues("name");
+		final List<String> ages = context.getStringValues("age");
+		assertNotNull(names);
+		assertNotNull(ages);
+		assertEquals(2, names.size());
+		assertEquals(2, ages.size());
+		assertEquals("Alice", names.get(0));
+		assertEquals("Bob", names.get(1));
+		assertEquals("25", ages.get(0));
+		assertEquals("30", ages.get(1));
+	}
+	
+	@Test
+	public void testSanitizeMissingClosingQuote() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		// With lenient JSON parsing enabled, this should work
+		final String jsonResponse = """
+				{
+					"field": "this is not right
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertEquals("this is not right", context.getStringValue("field"));
+	}
+	
+	@Test
+	public void testSanitizeUnescapedQuotesInString() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"field": '"yipii" said the clown'
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		final String value = context.getStringValue("field");
+		assertNotNull(value);
+		assertTrue(value.contains("yipii"));
+		assertTrue(value.contains("clown"));
+	}
+
+    @Test
+	public void testSanitizeUnescapedQuotesInString2() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"field": ""yipii" said the clown"
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		final String value = context.getStringValue("field");
+		assertNotNull(value);
+		assertTrue(value.contains("yipii"));
+		assertTrue(value.contains("clown"));
+	}
+	
+	@Test
+	public void testLenientJsonFeatures() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"greeting": 'Hello "World"',
+					"message": 'She said "hi" to me',
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertNotNull(context.getStringValue("greeting"));
+		assertNotNull(context.getStringValue("message"));
+	}
+	
+	@Test
+	public void testJsonWithEscapedCharacters() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"path": "C:\\\\Users\\\\Documents",
+					"quote": "He said \\"hello\\""
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertNotNull(context.getStringValue("path"));
+		assertNotNull(context.getStringValue("quote"));
+	}
+	
+	@Test
+	public void testInvalidJsonCannotBeSanitized() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					this is completely broken [[[
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertFalse(result);
+		
+		// Verify error message was added to conversation
+		final List<Conversation.Turn> turns = context.getConversation().getConversation();
+		final Conversation.Turn lastTurn = turns.get(turns.size() - 1);
+		assertTrue(lastTurn.text().contains("valid JSON syntax"));
+	}
+	
+	@Test
+	public void testJsonWithNoCurlyBraces() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = "This is just plain text with no JSON";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertFalse(result);
+	}
+	
+	@Test
+	public void testJsonWithAllowedFeatures() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		// Test various lenient JSON features that should be accepted
+		final String jsonResponse = """
+				{
+					// This is a comment
+					field1: 'single quotes',
+					field2: "trailing comma",
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertEquals("single quotes", context.getStringValue("field1"));
+		assertEquals("trailing comma", context.getStringValue("field2"));
+	}
+	
+	@Test
+	public void testJsonWithNumbers() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = """
+				{
+					"count": 42,
+					"price": 19.99,
+					"active": true
+				}
+				""";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertEquals("42", context.getStringValue("count"));
+		assertEquals("19.99", context.getStringValue("price"));
+		assertEquals("true", context.getStringValue("active"));
+	}
+	
+	@Test
+	public void testEmptyJson() {
+		final JsonAgentResponseHandler handler = new JsonAgentResponseHandler();
+		final AgentContext context = new AgentContext();
+		context.setConversation(Conversation.create("test"));
+		
+		final String jsonResponse = "{}";
+		
+		final Completion completion = new Completion(jsonResponse, null, null, null);
+		final boolean result = handler.handle(completion, context);
+		
+		assertTrue(result);
+		assertTrue(context.getContext().isEmpty());
+	}
+	
+}