@ai-sdk/google/src/google-json-accumulator.ts

export type PartialArg = {
  jsonPath: string;
  stringValue?: string | null;
  numberValue?: number | null;
  boolValue?: boolean | null;
  nullValue?: unknown;
  willContinue?: boolean | null;
};

type PathSegment = string | number;

type StackEntry = {
  segment: PathSegment;
  isArray: boolean;
  childCount: number;
};

/**
 * Incrementally builds a JSON object from Google's streaming `partialArgs`
 * chunks emitted during tool-call function calling. Tracks both the structured
 * object and a running JSON text representation so callers can emit text deltas
 * that, when concatenated, form valid nested JSON matching JSON.stringify output.
 *
 * Input: [{jsonPath:"$.location",stringValue:"Boston"}]
 * Output: '{"location":"Boston"', then finalize() → closingDelta='}'
 */
export class GoogleJSONAccumulator {
  private accumulatedArgs: Record<string, unknown> = {};
  private jsonText = '';

  /**
   * Stack representing the currently "open" containers in the JSON output.
   * Entry 0 is always the root `{` object once the first value is written.
   */
  private pathStack: StackEntry[] = [];

  /**
   * Whether a string value is currently "open" (willContinue was true),
   * meaning the closing quote has not yet been emitted.
   */
  private stringOpen = false;

  /**
   * Input: [{jsonPath:"$.brightness",numberValue:50}]
   * Output: { currentJSON:{brightness:50}, textDelta:'{"brightness":50' }
   */
  processPartialArgs(partialArgs: PartialArg[]): {
    currentJSON: Record<string, unknown>;
    textDelta: string;
  } {
    let delta = '';

    for (const arg of partialArgs) {
      // Vertex terminates an array with a value-less `$.path[]` entry. For an
      // empty array, that terminator is the only entry sent for the path.
      const isArrayTerminator = arg.jsonPath.endsWith('[]');
      const jsonPath = isArrayTerminator
        ? arg.jsonPath.slice(0, -2)
        : arg.jsonPath;
      const rawPath = jsonPath.replace(/^\$\./, '');
      if (!rawPath || rawPath === '$') continue;

      const segments = parsePath(rawPath);
      if (!normalizeAndCheckIndexes(this.accumulatedArgs, segments)) continue;

      const existingValue = getNestedValue(this.accumulatedArgs, segments);
      const isStringContinuation =
        arg.stringValue != null && existingValue !== undefined;

      if (isStringContinuation) {
        const escaped = JSON.stringify(arg.stringValue).slice(1, -1);
        setNestedValue(
          this.accumulatedArgs,
          segments,
          (existingValue as string) + arg.stringValue,
        );
        delta += escaped;
        if (arg.willContinue !== true) {
          delta += this.closeOpenString();
        }
        continue;
      }

      const resolved = resolvePartialArgValue(arg);
      if (resolved == null) {
        // A value-less entry for a path without a value is an empty container
        // (an empty container has no leaves to stream). Value-less entries for
        // existing paths are string or array terminators. A string terminator
        // closes the open string, an array terminator is ignored.
        if (existingValue !== undefined) {
          if (typeof existingValue === 'string' && arg.willContinue !== true) {
            delta += this.closeOpenString();
          }
          continue;
        }
        if (arg.willContinue === true) continue;

        setNestedValue(
          this.accumulatedArgs,
          segments,
          isArrayTerminator ? [] : {},
        );
        delta += this.emitContainer(segments, isArrayTerminator);
        continue;
      }

      setNestedValue(this.accumulatedArgs, segments, resolved.value);
      delta += this.emitNavigationTo(segments, arg, resolved.json);
    }

    this.jsonText += delta;

    return {
      currentJSON: this.accumulatedArgs,
      textDelta: delta,
    };
  }

  /**
   * Input: jsonText='{"brightness":50', accumulatedArgs={brightness:50}
   * Output: { finalJSON:'{"brightness":50}', closingDelta:'}' }
   */
  finalize(): { finalJSON: string; closingDelta: string } {
    const finalArgs = JSON.stringify(this.accumulatedArgs);
    const closingDelta = finalArgs.slice(this.jsonText.length);
    return { finalJSON: finalArgs, closingDelta };
  }

  /**
   * Input: stringOpen=true
   * Output: '"' (and stringOpen=false), or '' when no string is open
   */
  private closeOpenString(): string {
    if (!this.stringOpen) {
      return '';
    }
    this.stringOpen = false;
    return '"';
  }

  /**
   * Input: pathStack=[] (first call) or pathStack=[root,...] (subsequent calls)
   * Output: '{' (first call) or '' (subsequent calls)
   */
  private ensureRoot(): string {
    if (this.pathStack.length === 0) {
      this.pathStack.push({ segment: '', isArray: false, childCount: 0 });
      return '{';
    }
    return '';
  }

  /**
   * Emits the JSON text fragment needed to navigate from the current open
   * path to the new leaf at `targetSegments`, then writes the value.
   *
   * Input: targetSegments=["recipe","name"], arg={jsonPath:"$.recipe.name",stringValue:"Lasagna"}, valueJson='"Lasagna"'
   * Output: '{"recipe":{"name":"Lasagna"'
   */
  private emitNavigationTo(
    targetSegments: PathSegment[],
    arg: PartialArg,
    valueJson: string,
  ): string {
    let fragment = this.closeOpenString();

    fragment += this.ensureRoot();

    const targetContainerSegments = targetSegments.slice(0, -1);
    const leafSegment = targetSegments[targetSegments.length - 1];

    const commonDepth = this.findCommonStackDepth(targetContainerSegments);

    fragment += this.closeDownTo(commonDepth);
    fragment += this.openDownTo(targetContainerSegments, leafSegment);
    fragment += this.emitLeaf(leafSegment, arg, valueJson);

    return fragment;
  }

  /**
   * Navigates to `targetSegments` and opens a container there. The container
   * stays open so later entries for its children continue inside it; it is
   * closed by subsequent navigation or by `finalize()`.
   *
   * Input: targetSegments=["nested","tags"], isArray=true
   * Output: '{"nested":{"tags":['
   */
  private emitContainer(
    targetSegments: PathSegment[],
    isArray: boolean,
  ): string {
    let fragment = this.closeOpenString();

    fragment += this.ensureRoot();

    const commonDepth = this.findCommonStackDepth(targetSegments);

    fragment += this.closeDownTo(commonDepth);
    // openDownTo infers the innermost container type from the leaf segment type
    fragment += this.openDownTo(targetSegments, isArray ? 0 : '');

    return fragment;
  }

  /**
   * Returns the stack depth to preserve when navigating to a new target
   * container path. Always >= 1 (the root is never popped).
   *
   * Input: stack=[root,"recipe","ingredients",0], target=["recipe","ingredients",1]
   * Output: 3 (keep root+"recipe"+"ingredients")
   */
  private findCommonStackDepth(targetContainer: PathSegment[]): number {
    const maxDepth = Math.min(
      this.pathStack.length - 1,
      targetContainer.length,
    );
    let common = 0;
    for (let i = 0; i < maxDepth; i++) {
      if (this.pathStack[i + 1].segment === targetContainer[i]) {
        common++;
      } else {
        break;
      }
    }
    return common + 1;
  }

  /**
   * Closes containers from the current stack depth back down to `targetDepth`.
   *
   * Input: this.pathStack=[root,"recipe","ingredients",0], targetDepth=3
   * Output: '}'
   */
  private closeDownTo(targetDepth: number): string {
    let fragment = '';
    while (this.pathStack.length > targetDepth) {
      const entry = this.pathStack.pop()!;
      fragment += entry.isArray ? ']' : '}';
    }
    return fragment;
  }

  /**
   * Opens containers from the current stack depth down to the full target
   * container path, emitting opening `{`, `[`, keys, and commas as needed.
   * `leafSegment` is used to determine if the innermost container is an array.
   *
   * Input: this.pathStack=[root], targetContainer=["recipe","ingredients"], leafSegment=0
   * Output: '"recipe":{"ingredients":['
   */
  private openDownTo(
    targetContainer: PathSegment[],
    leafSegment: PathSegment,
  ): string {
    let fragment = '';

    const startIdx = this.pathStack.length - 1;

    for (let i = startIdx; i < targetContainer.length; i++) {
      const pathSegment = targetContainer[i];
      const parentEntry = this.pathStack[this.pathStack.length - 1];

      if (parentEntry.childCount > 0) {
        fragment += ',';
      }
      parentEntry.childCount++;

      if (typeof pathSegment === 'string') {
        fragment += `${JSON.stringify(pathSegment)}:`;
      }

      const childSeg =
        i + 1 < targetContainer.length ? targetContainer[i + 1] : leafSegment;
      const isArray = typeof childSeg === 'number';

      fragment += isArray ? '[' : '{';

      this.pathStack.push({ segment: pathSegment, isArray, childCount: 0 });
    }

    return fragment;
  }

  /**
   * Emits the comma, key, and value for a leaf entry in the current container.
   *
   * Input: leafSegment="name", arg={stringValue:"Lasagna"}, valueJson='"Lasagna"'
   * Output: '"name":"Lasagna"' (or ',"name":"Lasagna"' if container.childCount > 0)
   */
  private emitLeaf(
    leafSegment: PathSegment,
    arg: PartialArg,
    valueJson: string,
  ): string {
    let fragment = '';
    const container = this.pathStack[this.pathStack.length - 1];

    if (container.childCount > 0) {
      fragment += ',';
    }
    container.childCount++;

    if (typeof leafSegment === 'string') {
      fragment += `${JSON.stringify(leafSegment)}:`;
    }

    if (arg.stringValue != null && arg.willContinue) {
      fragment += valueJson.slice(0, -1);
      this.stringOpen = true;
    } else {
      fragment += valueJson;
    }

    return fragment;
  }
}

/**
 * Splits a dotted/bracketed JSON path like `recipe.ingredients[0].name` into segments.
 *
 * Input: "recipe.ingredients[0].name"
 * Output: ["recipe", "ingredients", 0, "name"]
 */
function parsePath(rawPath: string): Array<string | number> {
  const segments: Array<string | number> = [];
  for (const part of rawPath.split('.')) {
    const bracketIdx = part.indexOf('[');
    if (bracketIdx === -1) {
      segments.push(part);
    } else {
      if (bracketIdx > 0) segments.push(part.slice(0, bracketIdx));
      for (const m of part.matchAll(/\[(\d+)\]/g)) {
        segments.push(parseInt(m[1], 10));
      }
    }
  }
  return segments;
}

const hasOwn = Object.prototype.hasOwnProperty;

/**
 * Checks only direct properties so path traversal never follows the prototype chain.
 */
function hasOwnProperty(
  obj: Record<string | number, unknown>,
  key: string | number,
): boolean {
  return hasOwn.call(obj, key);
}

/**
 * Defines path values as own data properties so special keys like `__proto__`
 * cannot invoke prototype setters while accumulating streamed arguments.
 */
function defineOwnProperty(
  obj: Record<string | number, unknown>,
  key: string | number,
  value: unknown,
): void {
  Object.defineProperty(obj, key, {
    value,
    enumerable: true,
    configurable: true,
    writable: true,
  });
}

/**
 * Traverses a nested object along the given path segments and returns the leaf value.
 *
 * Input: ({recipe:{name:"Lasagna"}}, ["recipe","name"])
 * Output: "Lasagna"
 */
function getNestedValue(
  obj: Record<string, unknown>,
  segments: Array<string | number>,
): unknown {
  let current: unknown = obj;
  for (const pathSegment of segments) {
    if (current == null || typeof current !== 'object') return undefined;
    const currentRecord = current as Record<string | number, unknown>;
    if (!hasOwnProperty(currentRecord, pathSegment)) return undefined;
    current = currentRecord[pathSegment];
  }
  return current;
}

/**
 * Normalizes dotted array indexes in place and checks that every array index
 * is an existing element or the next one. Arrays are streamed in order, and a
 * sparse index would make `finalize()` serialize one `null` per hole, so a tiny
 * response could allocate without bound.
 * A missing or `null` parent is replaced with a new array, so it starts at 0.
 *
 * Input: ({items:["a"]}, ["items",1]) or ({items:["a"]}, ["items",5])
 * Output: true or false
 */
function normalizeAndCheckIndexes(
  obj: Record<string, unknown>,
  segments: Array<string | number>,
): boolean {
  let current: unknown = obj;
  for (let i = 0; i < segments.length; i++) {
    let pathSegment = segments[i];
    if (Array.isArray(current) && typeof pathSegment === 'string') {
      const index = Number(pathSegment);
      // Only canonical array-index keys affect length. Keep object keys as strings.
      if (
        Number.isInteger(index) &&
        index >= 0 &&
        index < 2 ** 32 - 1 &&
        String(index) === pathSegment
      ) {
        pathSegment = index;
        segments[i] = index;
      }
    }
    if (typeof pathSegment === 'number') {
      const length = Array.isArray(current)
        ? current.length
        : current == null
          ? 0
          : undefined;
      if (length !== undefined && pathSegment > length) return false;
    }
    current =
      current != null &&
      typeof current === 'object' &&
      hasOwnProperty(current as Record<string | number, unknown>, pathSegment)
        ? (current as Record<string | number, unknown>)[pathSegment]
        : undefined;
  }
  return true;
}

/**
 * Sets a value at a nested path, creating intermediate objects or arrays as needed.
 *
 * Input: obj={}, segments=["recipe","ingredients",0,"name"], value="Noodles"
 * Output: {recipe:{ingredients:[{name:"Noodles"}]}}
 */
function setNestedValue(
  obj: Record<string, unknown>,
  segments: Array<string | number>,
  value: unknown,
): void {
  let current: Record<string | number, unknown> = obj;
  for (let i = 0; i < segments.length - 1; i++) {
    const pathSegment = segments[i];
    const nextSeg = segments[i + 1];
    if (!hasOwnProperty(current, pathSegment) || current[pathSegment] == null) {
      defineOwnProperty(
        current,
        pathSegment,
        typeof nextSeg === 'number' ? [] : {},
      );
    }
    current = current[pathSegment] as Record<string | number, unknown>;
  }
  defineOwnProperty(current, segments[segments.length - 1], value);
}

/**
 * Extracts the first non-null typed value from a partial arg and returns it with its JSON representation.
 *
 * Input: arg={stringValue:"Boston"} or arg={numberValue:50}
 * Output: {value:"Boston", json:'"Boston"'} or {value:50, json:'50'}
 */
function resolvePartialArgValue(arg: {
  stringValue?: string | null;
  numberValue?: number | null;
  boolValue?: boolean | null;
  nullValue?: unknown;
}): { value: unknown; json: string } | undefined {
  const value = arg.stringValue ?? arg.numberValue ?? arg.boolValue;
  if (value != null) return { value, json: JSON.stringify(value) };
  if ('nullValue' in arg) return { value: null, json: 'null' };
}