Jyotisha/frontend/src/lib/stream-frame-buffer.ts

/**
 * Frame-coalesced release of streamed agent output.
 *
 * Every network chunk used to become its own React commit, and each commit
 * re-parsed the whole partial answer. This buffer sits between the event
 * parser and `setState`: events mutate an accumulator, and at most one flush
 * happens per animation frame. Answer and thinking text are released at a
 * steady per-frame pace so a burst of chunks reads as flowing text instead of
 * a jump, while a large backlog (reconnect, slow tab) catches up in roughly a
 * dozen frames. When the stream ends, what is still unreleased is written out
 * at the same pace within a short budget instead of appearing in one frame
 * (BUG-1075); only a stop, a failure or a hidden document releases at once.
 *
 * Pure release arithmetic lives in exported functions so the policy is
 * testable without a DOM; scheduling is injectable for the same reason.
 */

export const STREAM_RELEASE_MIN_CHARS = 2;
/**
 * Typing pace: at most this many characters per frame while the backlog is
 * ordinary (BUG-1075). Four a frame is about 240 characters a second at 60 fps,
 * above what a model produces on average, so the backlog does not grow; a
 * server-held lump (the 160-character answer release, a whole short follow-up)
 * reads as writing instead of appearing at once.
 */
export const STREAM_RELEASE_MAX_CHARS = 4;
export const STREAM_RELEASE_CATCHUP_DIVISOR = 12;
/** A backlog above this (reconnect, slow tab) catches up in about twelve frames instead of typing it out. */
export const STREAM_RELEASE_CATCHUP_CHARS = 600;
/** A paced settle writes out what is left within this many frames (1.5 s at 60 fps). */
export const STREAM_SETTLE_MAX_FRAMES = 90;
export const STREAM_HIDDEN_FLUSH_MS = 250;

/**
 * Characters to reveal on one frame. `backlogChars` is how much was waiting
 * when the newest text arrived. An ordinary backlog is typed out at up to
 * STREAM_RELEASE_MAX_CHARS a frame, with a two-character floor so a slow model
 * never reads as stalled; a backlog above STREAM_RELEASE_CATCHUP_CHARS is
 * divided by twelve so any burst clears in about twelve frames. Callers
 * without a backlog figure pass the pending count.
 */
export function streamReleaseCount(pendingChars: number, backlogChars = pendingChars): number {
  if (pendingChars <= 0) return 0;
  const catchUp = Math.ceil(backlogChars / STREAM_RELEASE_CATCHUP_DIVISOR);
  const perFrame = backlogChars > STREAM_RELEASE_CATCHUP_CHARS
    ? catchUp
    : Math.min(STREAM_RELEASE_MAX_CHARS, catchUp);
  return Math.min(pendingChars, Math.max(STREAM_RELEASE_MIN_CHARS, perFrame));
}

/** Characters per frame that write out `pendingChars` within the settle budget, never slower than the typing pace. */
export function settleReleaseCount(pendingChars: number): number {
  if (pendingChars <= 0) return 0;
  return Math.max(STREAM_RELEASE_MAX_CHARS, Math.ceil(pendingChars / STREAM_SETTLE_MAX_FRAMES));
}

/** Advance a released prefix toward its target by one frame's worth of text. */
export function advanceStreamRelease(released: string, target: string, backlogChars?: number, perFrame?: number): string {
  if (!target.startsWith(released)) {
    // The target was replaced rather than extended: restart from its head.
    return target.slice(0, perFrame ?? streamReleaseCount(target.length, backlogChars ?? target.length));
  }
  const pending = target.length - released.length;
  if (pending <= 0) return target;
  return target.slice(0, released.length + (perFrame ?? streamReleaseCount(pending, backlogChars ?? pending)));
}

export type StreamFrameSnapshot<Meta> = Readonly<{
  answer: string;
  thinking: string;
  meta: Meta;
  /** True when this flush released everything that had arrived. */
  settled: boolean;
}>;

export type StreamFrameScheduler = Readonly<{
  requestFrame: (callback: () => void) => number;
  cancelFrame: (handle: number) => void;
  requestTimeout: (callback: () => void, delayMs: number) => number;
  cancelTimeout: (handle: number) => void;
  hidden: () => boolean;
}>;

export type StreamFrameBufferOptions<Meta> = Readonly<{
  initialMeta: Meta;
  flush: (snapshot: StreamFrameSnapshot<Meta>) => void;
  scheduler?: StreamFrameScheduler;
}>;

export type StreamFrameBuffer<Meta> = Readonly<{
  setAnswer: (fullText: string) => void;
  setThinking: (fullText: string) => void;
  setMeta: (next: Meta | ((current: Meta) => Meta)) => void;
  /** Publish meta-only changes (timeline rows, activity) on the next frame. */
  touch: () => void;
  /**
   * The stream has ended. By default what is still unreleased is written out
   * at the typing pace within STREAM_SETTLE_MAX_FRAMES, and the last frame
   * flushes with `settled: true`; the promise resolves on that frame. With
   * `immediate` (stop, failure, disconnect) everything is released in one
   * synchronous flush, as before BUG-1075. A hidden document always releases
   * at once.
   */
  settle: (options?: Readonly<{ immediate?: boolean }>) => Promise<void>;
  /** Drop everything, including scheduled work, without flushing. */
  reset: (meta?: Meta) => void;
  dispose: () => void;
  /** Text released so far, for callers that persist partial output. */
  released: () => Readonly<{ answer: string; thinking: string }>;
}>;

function pendingChars(released: string, target: string): number {
  return target.startsWith(released) ? target.length - released.length : target.length;
}

function browserScheduler(): StreamFrameScheduler {
  return {
    requestFrame: (callback) => window.requestAnimationFrame(callback),
    cancelFrame: (handle) => window.cancelAnimationFrame(handle),
    requestTimeout: (callback, delayMs) => window.setTimeout(callback, delayMs),
    cancelTimeout: (handle) => window.clearTimeout(handle),
    hidden: () => typeof document !== "undefined" && document.hidden,
  };
}

export function createStreamFrameBuffer<Meta>(
  options: StreamFrameBufferOptions<Meta>,
): StreamFrameBuffer<Meta> {
  const scheduler = options.scheduler ?? browserScheduler();
  let targetAnswer = "";
  let targetThinking = "";
  let releasedAnswer = "";
  let releasedThinking = "";
  let answerBacklog = 0;
  let thinkingBacklog = 0;
  let meta = options.initialMeta;
  let disposed = false;
  let frameHandle: number | null = null;
  let timeoutHandle: number | null = null;
  // A paced settle in flight: its per-frame count, and who to tell when the
  // last frame has flushed. dispose() during it defers until that frame.
  let settling: { perFrame: number; resolve: () => void } | null = null;
  let disposeWhenSettled = false;

  const cancelScheduled = () => {
    if (frameHandle !== null) {
      scheduler.cancelFrame(frameHandle);
      frameHandle = null;
    }
    if (timeoutHandle !== null) {
      scheduler.cancelTimeout(timeoutHandle);
      timeoutHandle = null;
    }
  };

  const emit = (settled: boolean) => {
    options.flush({
      answer: releasedAnswer,
      thinking: releasedThinking,
      meta,
      settled,
    });
  };

  const finishSettle = () => {
    const done = settling;
    settling = null;
    if (disposeWhenSettled) {
      disposeWhenSettled = false;
      disposed = true;
    }
    done?.resolve();
  };

  const step = () => {
    frameHandle = null;
    timeoutHandle = null;
    if (disposed) return;
    if (scheduler.hidden()) {
      releasedAnswer = targetAnswer;
      releasedThinking = targetThinking;
    } else if (settling) {
      releasedAnswer = advanceStreamRelease(releasedAnswer, targetAnswer, answerBacklog, settling.perFrame);
      releasedThinking = targetThinking;
    } else {
      releasedAnswer = advanceStreamRelease(releasedAnswer, targetAnswer, answerBacklog);
      releasedThinking = advanceStreamRelease(releasedThinking, targetThinking, thinkingBacklog);
    }
    if (releasedAnswer === targetAnswer) answerBacklog = 0;
    if (releasedThinking === targetThinking) thinkingBacklog = 0;
    const caughtUp = releasedAnswer === targetAnswer && releasedThinking === targetThinking;
    emit(caughtUp);
    if (!caughtUp) {
      schedule();
      return;
    }
    if (settling) finishSettle();
  };

  const schedule = () => {
    if (disposed || frameHandle !== null || timeoutHandle !== null) return;
    if (scheduler.hidden()) {
      timeoutHandle = scheduler.requestTimeout(step, STREAM_HIDDEN_FLUSH_MS);
    } else {
      frameHandle = scheduler.requestFrame(step);
    }
  };

  return {
    setAnswer(fullText) {
      if (disposed || fullText === targetAnswer) return;
      targetAnswer = fullText;
      answerBacklog = Math.max(answerBacklog, pendingChars(releasedAnswer, targetAnswer));
      schedule();
    },
    setThinking(fullText) {
      if (disposed || fullText === targetThinking) return;
      targetThinking = fullText;
      thinkingBacklog = Math.max(thinkingBacklog, pendingChars(releasedThinking, targetThinking));
      schedule();
    },
    setMeta(next) {
      if (disposed) return;
      meta = typeof next === "function" ? (next as (current: Meta) => Meta)(meta) : next;
      schedule();
    },
    touch() {
      if (disposed) return;
      schedule();
    },
    settle(settleOptions) {
      if (disposed) return Promise.resolve();
      if (settling) {
        // A second settle while one is writing out: an immediate one takes
        // over and flushes now; a paced one just waits for the first.
        if (!settleOptions?.immediate) {
          return new Promise<void>((resolve) => {
            const previous = settling!.resolve;
            settling!.resolve = () => { previous(); resolve(); };
          });
        }
        cancelScheduled();
      }
      const pending = pendingChars(releasedAnswer, targetAnswer);
      if (settleOptions?.immediate || scheduler.hidden() || pending <= 0) {
        cancelScheduled();
        releasedAnswer = targetAnswer;
        releasedThinking = targetThinking;
        answerBacklog = 0;
        thinkingBacklog = 0;
        emit(true);
        if (settling) finishSettle();
        return Promise.resolve();
      }
      return new Promise<void>((resolve) => {
        settling = { perFrame: settleReleaseCount(pending), resolve };
        schedule();
      });
    },
    reset(nextMeta) {
      cancelScheduled();
      if (settling) finishSettle();
      targetAnswer = "";
      targetThinking = "";
      releasedAnswer = "";
      releasedThinking = "";
      answerBacklog = 0;
      thinkingBacklog = 0;
      if (nextMeta !== undefined) meta = nextMeta;
    },
    dispose() {
      if (settling) {
        // Let the paced settle write out its last frames; it disposes itself.
        disposeWhenSettled = true;
        return;
      }
      disposed = true;
      cancelScheduled();
    },
    released() {
      return { answer: releasedAnswer, thinking: releasedThinking };
    },
  };
}