@assistant-ui/react/src/utils/json/fix-json.ts

Typescript/React library for AI Chat

// LICENSE for this file only

// Copyright 2023 Vercel, Inc.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

type State =
  | "ROOT"
  | "FINISH"
  | "INSIDE_STRING"
  | "INSIDE_STRING_ESCAPE"
  | "INSIDE_LITERAL"
  | "INSIDE_NUMBER"
  | "INSIDE_OBJECT_START"
  | "INSIDE_OBJECT_KEY"
  | "INSIDE_OBJECT_AFTER_KEY"
  | "INSIDE_OBJECT_BEFORE_VALUE"
  | "INSIDE_OBJECT_AFTER_VALUE"
  | "INSIDE_OBJECT_AFTER_COMMA"
  | "INSIDE_ARRAY_START"
  | "INSIDE_ARRAY_AFTER_VALUE"
  | "INSIDE_ARRAY_AFTER_COMMA";

// Implemented as a scanner with additional fixing
// that performs a single linear time scan pass over the partial JSON.
//
// The states should ideally match relevant states from the JSON spec:
// https://www.json.org/json-en.html
//
// Please note that invalid JSON is not considered/covered, because it
// is assumed that the resulting JSON will be processed by a standard
// JSON parser that will detect any invalid JSON.
export function fixJson(input: string): [string, number] {
  const stack: State[] = ["ROOT"];
  let lastValidIndex = -1;
  let literalStart: number | null = null;

  function processValueStart(char: string, i: number, swapState: State) {
    {
      switch (char) {
        case '"': {
          lastValidIndex = i;
          stack.pop();
          stack.push(swapState);
          stack.push("INSIDE_STRING");
          break;
        }

        case "f":
        case "t":
        case "n": {
          lastValidIndex = i;
          literalStart = i;
          stack.pop();
          stack.push(swapState);
          stack.push("INSIDE_LITERAL");
          break;
        }

        case "-": {
          stack.pop();
          stack.push(swapState);
          stack.push("INSIDE_NUMBER");
          break;
        }
        case "0":
        case "1":
        case "2":
        case "3":
        case "4":
        case "5":
        case "6":
        case "7":
        case "8":
        case "9": {
          lastValidIndex = i;
          stack.pop();
          stack.push(swapState);
          stack.push("INSIDE_NUMBER");
          break;
        }

        case "{": {
          lastValidIndex = i;
          stack.pop();
          stack.push(swapState);
          stack.push("INSIDE_OBJECT_START");
          break;
        }

        case "[": {
          lastValidIndex = i;
          stack.pop();
          stack.push(swapState);
          stack.push("INSIDE_ARRAY_START");
          break;
        }
      }
    }
  }

  function processAfterObjectValue(char: string, i: number) {
    switch (char) {
      case ",": {
        stack.pop();
        stack.push("INSIDE_OBJECT_AFTER_COMMA");
        break;
      }
      case "}": {
        lastValidIndex = i;
        stack.pop();
        break;
      }
    }
  }

  function processAfterArrayValue(char: string, i: number) {
    switch (char) {
      case ",": {
        stack.pop();
        stack.push("INSIDE_ARRAY_AFTER_COMMA");
        break;
      }
      case "]": {
        lastValidIndex = i;
        stack.pop();
        break;
      }
    }
  }

  for (let i = 0; i < input.length; i++) {
    const char = input[i]!;
    const currentState = stack[stack.length - 1];

    switch (currentState) {
      case "ROOT":
        processValueStart(char, i, "FINISH");
        break;

      case "INSIDE_OBJECT_START": {
        switch (char) {
          case '"': {
            stack.pop();
            stack.push("INSIDE_OBJECT_KEY");
            break;
          }
          case "}": {
            lastValidIndex = i;
            stack.pop();
            break;
          }
        }
        break;
      }

      case "INSIDE_OBJECT_AFTER_COMMA": {
        switch (char) {
          case '"': {
            stack.pop();
            stack.push("INSIDE_OBJECT_KEY");
            break;
          }
        }
        break;
      }

      case "INSIDE_OBJECT_KEY": {
        switch (char) {
          case '"': {
            stack.pop();
            stack.push("INSIDE_OBJECT_AFTER_KEY");
            break;
          }
        }
        break;
      }

      case "INSIDE_OBJECT_AFTER_KEY": {
        switch (char) {
          case ":": {
            stack.pop();
            stack.push("INSIDE_OBJECT_BEFORE_VALUE");

            break;
          }
        }
        break;
      }

      case "INSIDE_OBJECT_BEFORE_VALUE": {
        processValueStart(char, i, "INSIDE_OBJECT_AFTER_VALUE");
        break;
      }

      case "INSIDE_OBJECT_AFTER_VALUE": {
        processAfterObjectValue(char, i);
        break;
      }

      case "INSIDE_STRING": {
        switch (char) {
          case '"': {
            stack.pop();
            lastValidIndex = i;
            break;
          }

          case "\\": {
            stack.push("INSIDE_STRING_ESCAPE");
            break;
          }

          default: {
            lastValidIndex = i;
          }
        }

        break;
      }

      case "INSIDE_ARRAY_START": {
        switch (char) {
          case "]": {
            lastValidIndex = i;
            stack.pop();
            break;
          }

          default: {
            lastValidIndex = i;
            processValueStart(char, i, "INSIDE_ARRAY_AFTER_VALUE");
            break;
          }
        }
        break;
      }

      case "INSIDE_ARRAY_AFTER_VALUE": {
        switch (char) {
          case ",": {
            stack.pop();
            stack.push("INSIDE_ARRAY_AFTER_COMMA");
            break;
          }

          case "]": {
            lastValidIndex = i;
            stack.pop();
            break;
          }

          default: {
            lastValidIndex = i;
            break;
          }
        }

        break;
      }

      case "INSIDE_ARRAY_AFTER_COMMA": {
        processValueStart(char, i, "INSIDE_ARRAY_AFTER_VALUE");
        break;
      }

      case "INSIDE_STRING_ESCAPE": {
        stack.pop();
        lastValidIndex = i;

        break;
      }

      case "INSIDE_NUMBER": {
        switch (char) {
          case "0":
          case "1":
          case "2":
          case "3":
          case "4":
          case "5":
          case "6":
          case "7":
          case "8":
          case "9": {
            lastValidIndex = i;
            break;
          }

          case "e":
          case "E":
          case "-":
          case ".": {
            break;
          }

          case ",": {
            stack.pop();

            if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
              processAfterArrayValue(char, i);
            }

            if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
              processAfterObjectValue(char, i);
            }

            break;
          }

          case "}": {
            stack.pop();

            if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
              processAfterObjectValue(char, i);
            }

            break;
          }

          case "]": {
            stack.pop();

            if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
              processAfterArrayValue(char, i);
            }

            break;
          }

          default: {
            stack.pop();
            break;
          }
        }

        break;
      }

      case "INSIDE_LITERAL": {
        const partialLiteral = input.substring(literalStart!, i + 1);

        if (
          !"false".startsWith(partialLiteral) &&
          !"true".startsWith(partialLiteral) &&
          !"null".startsWith(partialLiteral)
        ) {
          stack.pop();

          if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
            processAfterObjectValue(char, i);
          } else if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
            processAfterArrayValue(char, i);
          }
        } else {
          lastValidIndex = i;
        }

        break;
      }
    }
  }

  let result = input.slice(0, lastValidIndex + 1);
  let partialCount = 0;

  for (let i = stack.length - 1; i >= 0; i--) {
    const state = stack[i];

    switch (state) {
      case "INSIDE_STRING": {
        result += '"';
        partialCount++;
        break;
      }

      case "INSIDE_OBJECT_KEY":
      case "INSIDE_OBJECT_AFTER_KEY":
      case "INSIDE_OBJECT_AFTER_COMMA":
      case "INSIDE_OBJECT_START":
      case "INSIDE_OBJECT_BEFORE_VALUE":
      case "INSIDE_OBJECT_AFTER_VALUE": {
        result += "}";
        partialCount++;
        break;
      }

      case "INSIDE_ARRAY_START":
      case "INSIDE_ARRAY_AFTER_COMMA":
      case "INSIDE_ARRAY_AFTER_VALUE": {
        result += "]";
        partialCount++;
        break;
      }

      case "INSIDE_LITERAL": {
        const partialLiteral = input.substring(literalStart!, input.length);

        if ("true".startsWith(partialLiteral)) {
          result += "true".slice(partialLiteral.length);
        } else if ("false".startsWith(partialLiteral)) {
          result += "false".slice(partialLiteral.length);
        } else if ("null".startsWith(partialLiteral)) {
          result += "null".slice(partialLiteral.length);
        }
      }
    }
  }

  return [result, partialCount];
}