All files / kernel-agents/src/strategies/json sample-collector.ts

94.44% Statements 17/18
90% Branches 9/10
100% Functions 2/2
94.44% Lines 17/18

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61                                          2x                 11x 11x 11x 16x 16x 16x   16x 16x 16x 16x   8x   2x 2x           6x 1x         5x      
import { SampleGenerationError } from '@metamask/kernel-errors';
import type { Logger } from '@metamask/logger';
import { parse, MalformedJSON, PartialJSON } from 'partial-json';
 
import type { SampleCollector } from '../../types.ts';
 
/**
 * A sample collector for a streaming JSON object response.
 *
 * Uses `partial-json` for a three-way check on each chunk:
 * - **Complete** — `parse(buffer, 0)` returns the parsed object.
 * - **Incomplete** — throws `PartialJSON`, meaning the buffer is valid so far.
 * - **Malformed** — throws `MalformedJSON`, meaning the buffer is irrecoverable.
 *
 * @param args - The arguments to make the sample collector.
 * @param args.prefix - The prefix to prepend to the response
 * @param args.maxChunkCount - The maximum number of chunks to parse
 * @param args.logger - The logger to use for the sample collector
 * @returns A function that collects a delta of a streaming response,
 *   returning the result value if collecting is complete or null otherwise.
 */
export const makeSampleCollector = <Result = unknown>({
  prefix = '',
  maxChunkCount = 200,
  logger,
}: {
  prefix?: string;
  maxChunkCount?: number;
  logger?: Logger;
}): SampleCollector<Result> => {
  let response = prefix;
  let chunkCount = 0;
  return (delta: string) => {
    chunkCount += 1;
    response += delta;
    logger?.info('toParse:', response);
 
    try {
      const result = parse(response, 0);
      logger?.info('parsed:', result);
      return result;
    } catch (error) {
      if (error instanceof PartialJSON) {
        // Buffer is incomplete but structurally valid.
      } else if (error instanceof MalformedJSON) {
        throw new SampleGenerationError(response, error);
      } else E{
        throw error;
      }
    }
 
    if (maxChunkCount && chunkCount > maxChunkCount) {
      throw new SampleGenerationError(
        response,
        new Error('Max chunk count reached'),
      );
    }
    return null;
  };
};