Fix an incomplete JSON to easily stream AI json responses

JsonFixer.java

package com.example.demo;

import java.util.ArrayDeque;
import java.util.Deque;

public class JsonFixer {

    public String fix(String response) {
        String json = findJson(response);
        if (json == null || json.isEmpty()) {
            return "[]";
        }
        json = json.trim().replaceAll("\n", "");
        // Try to parse and fix truncated JSON
        StringBuilder result = new StringBuilder();
        Deque<Character> stack = new ArrayDeque<>(); // tracks open [ and {
        boolean inString = false;
        boolean escaped = false;
        boolean lastCharWasSpace = false;
        // Track position where last comma was appended (for truncation cleanup in arrays)
        int lastArrayCommaPos = -1;
        // Track object state: after key, after colon, after value
        // States within an object:
        //  EXPECT_KEY, IN_KEY, AFTER_KEY, AFTER_COLON, IN_VALUE, AFTER_VALUE
        enum ObjState { EXPECT_KEY, IN_KEY, AFTER_KEY, AFTER_COLON, IN_VALUE, AFTER_VALUE }
        Deque<ObjState> objStateStack = new ArrayDeque<>();

        for (int i = 0; i < json.length(); i++) {
            char c = json.charAt(i);

            if (escaped) {
                result.append(c);
                escaped = false;
                continue;
            }

            if (inString) {
                if (c == '\\') {
                    escaped = true;
                    result.append(c);
                } else if (c == '"') {
                    inString = false;
                    result.append(c);
                    // Update object state
                    if (!objStateStack.isEmpty()) {
                        ObjState state = objStateStack.peek();
                        if (state == ObjState.IN_KEY) {
                            objStateStack.pop();
                            objStateStack.push(ObjState.AFTER_KEY);
                        } else if (state == ObjState.IN_VALUE) {
                            objStateStack.pop();
                            objStateStack.push(ObjState.AFTER_VALUE);
                        }
                    }
                } else {
                    result.append(c);
                }
                lastCharWasSpace = false;
                continue;
            }

            // Not in string — collapse whitespace outside strings
            if (c == ' ' || c == '\t') {
                if (!lastCharWasSpace && result.length() > 0) {
                    result.append(' ');
                    lastCharWasSpace = true;
                }
                continue;
            }
            lastCharWasSpace = false;

            if (c == '"') {
                inString = true;
                result.append(c);
                // Update object state
                if (!objStateStack.isEmpty()) {
                    ObjState state = objStateStack.peek();
                    if (state == ObjState.EXPECT_KEY) {
                        objStateStack.pop();
                        objStateStack.push(ObjState.IN_KEY);
                    } else if (state == ObjState.AFTER_COLON) {
                        objStateStack.pop();
                        objStateStack.push(ObjState.IN_VALUE);
                    }
                }
            } else if (c == '{') {
                result.append(c);
                stack.push('{');
                objStateStack.push(ObjState.EXPECT_KEY);
            } else if (c == '[') {
                result.append(c);
                stack.push('[');
            } else if (c == '}') {
                result.append(c);
                if (!stack.isEmpty() && stack.peek() == '{') {
                    stack.pop();
                }
                if (!objStateStack.isEmpty()) {
                    objStateStack.pop();
                }
                // Update parent object state if this was a value
                if (!objStateStack.isEmpty()) {
                    ObjState state = objStateStack.peek();
                    if (state == ObjState.AFTER_COLON || state == ObjState.IN_VALUE) {
                        objStateStack.pop();
                        objStateStack.push(ObjState.AFTER_VALUE);
                    }
                }
            } else if (c == ']') {
                result.append(c);
                if (!stack.isEmpty() && stack.peek() == '[') {
                    stack.pop();
                }
            } else if (c == ':') {
                result.append(c);
                if (!objStateStack.isEmpty()) {
                    ObjState state = objStateStack.peek();
                    if (state == ObjState.AFTER_KEY) {
                        objStateStack.pop();
                        objStateStack.push(ObjState.AFTER_COLON);
                    }
                }
            } else if (c == ',') {
                result.append(c);
                // Track array-level comma position for truncation cleanup
                if (!stack.isEmpty() && stack.peek() == '[') {
                    lastArrayCommaPos = result.length() - 1;
                }
                if (!objStateStack.isEmpty()) {
                    ObjState state = objStateStack.peek();
                    // Only transition to EXPECT_KEY when inside an object (not an array)
                    if (state == ObjState.AFTER_VALUE && !stack.isEmpty() && stack.peek() == '{') {
                        objStateStack.pop();
                        objStateStack.push(ObjState.EXPECT_KEY);
                    }
                }
            } else {
                result.append(c);
            }
        }

        // Now fix the truncated parts
        // If we're still in a string and we're inside an array (not object key/value),
        // this is likely a truncated array element — remove it back to the last comma
        if (inString && !stack.isEmpty() && stack.peek() == '[' && lastArrayCommaPos >= 0) {
            // Check if the object state is NOT in a key/value context for the current array
            boolean inArrayDirectly = objStateStack.isEmpty()
                || objStateStack.peek() == ObjState.AFTER_VALUE
                || objStateStack.peek() == ObjState.AFTER_COLON;
            if (inArrayDirectly) {
                // Remove everything from the last array comma
                result.setLength(lastArrayCommaPos);
            } else {
                // Close the string normally
                result.append('"');
                if (!objStateStack.isEmpty()) {
                    ObjState state = objStateStack.peek();
                    if (state == ObjState.IN_KEY) {
                        objStateStack.pop();
                        objStateStack.push(ObjState.AFTER_KEY);
                    } else if (state == ObjState.IN_VALUE) {
                        objStateStack.pop();
                        objStateStack.push(ObjState.AFTER_VALUE);
                    }
                }
            }
            inString = false;
        } else if (inString) {
            result.append('"');
            // Update state
            if (!objStateStack.isEmpty()) {
                ObjState state = objStateStack.peek();
                if (state == ObjState.IN_KEY) {
                    objStateStack.pop();
                    objStateStack.push(ObjState.AFTER_KEY);
                } else if (state == ObjState.IN_VALUE) {
                    objStateStack.pop();
                    objStateStack.push(ObjState.AFTER_VALUE);
                }
            }
        }

        // Handle incomplete object states
        if (!objStateStack.isEmpty()) {
            ObjState state = objStateStack.peek();
            if (state == ObjState.AFTER_KEY) {
                // Key complete but no colon/value — add empty value
                result.append(": \"\"");
                objStateStack.pop();
                objStateStack.push(ObjState.AFTER_VALUE);
            } else if (state == ObjState.AFTER_COLON) {
                // Colon present but no value — add empty value
                result.append(" \"\"");
                objStateStack.pop();
                objStateStack.push(ObjState.AFTER_VALUE);
            } else if (state == ObjState.EXPECT_KEY) {
                // Trailing comma — remove it
                String s = result.toString();
                if (s.endsWith(",")) {
                    result.setLength(result.length() - 1);
                } else if (s.endsWith(", ")) {
                    result.setLength(result.length() - 2);
                }
                objStateStack.pop();
                objStateStack.push(ObjState.AFTER_VALUE);
            }
        }


        // Close open containers
        while (!stack.isEmpty()) {
            char open = stack.pop();
            if (open == '{') {
                result.append('}');
            } else if (open == '[') {
                // Remove trailing comma before closing array
                String s = result.toString();
                if (s.endsWith(",")) {
                    result.setLength(result.length() - 1);
                } else if (s.endsWith(", ")) {
                    result.setLength(result.length() - 2);
                }
                result.append(']');
            }
        }

        return result.toString();
    }

    private String findJson(String response) {
        int jsonStart = response.indexOf("[");
        if (jsonStart >= 0) {
            int lastGraph = response.lastIndexOf("}");
            int lastSquare = response.lastIndexOf("]");
            int end = Math.max(lastGraph, lastSquare);
            return response.substring(jsonStart, end > 0 ? end : response.length());
        } else {
            return "";
        }
    }
}