next-tjwater-agent/src/routes/chatTokenSmoother.ts

type TokenSmootherOptions = {
  enabled: boolean;
  delayMs: number;
  locale: string;
  maxWaitMs?: number;
  sessionId: string;
  write: (event: string, data: Record<string, unknown>) => void;
};

const segmenters = new Map<string, Intl.Segmenter | null>();
const CJK_SCRIPT_PATTERN = /[\p{Script=Han}\p{Script=Hiragana}\p{Script=Katakana}\p{Script=Hangul}]/u;
const CJK_PUNCTUATION_PATTERN = /[。！？!?，、；：,;:]/u;
const CJK_FALLBACK_TRIGGER_CHARS = 24;
const CJK_FALLBACK_MIN_CHARS = 12;
const CJK_FALLBACK_TARGET_CHARS = 18;

function getTokenSmoothingSegmenter(locale: string) {
  if (segmenters.has(locale)) {
    return segmenters.get(locale) ?? null;
  }

  const segmenter =
    typeof Intl !== "undefined" && "Segmenter" in Intl
      ? new Intl.Segmenter(locale, { granularity: "word" })
      : null;
  segmenters.set(locale, segmenter);
  return segmenter;
}

function sliceCodePoints(value: string, count: number) {
  return Array.from(value).slice(0, count).join("");
}

function codePointLength(value: string) {
  return Array.from(value).length;
}

function isCjkSmoothingContent(content: string, locale: string) {
  return (
    CJK_SCRIPT_PATTERN.test(content) ||
    (/^(zh|ja|ko)(?:-|$)/i.test(locale) && CJK_PUNCTUATION_PATTERN.test(content))
  );
}

function detectCjkFallbackChunk(content: string, locale: string) {
  if (codePointLength(content) < CJK_FALLBACK_TRIGGER_CHARS) {
    return null;
  }

  const segmenter = getTokenSmoothingSegmenter(locale);
  if (segmenter) {
    let candidate: string | null = null;
    for (const segment of segmenter.segment(content)) {
      const nextCandidate = content.slice(0, segment.index + segment.segment.length);
      const nextLength = codePointLength(nextCandidate);
      if (nextLength > CJK_FALLBACK_TARGET_CHARS) {
        break;
      }
      if (nextLength >= CJK_FALLBACK_MIN_CHARS) {
        candidate = nextCandidate;
      }
    }
    if (candidate) {
      return candidate;
    }
  }

  return sliceCodePoints(content, CJK_FALLBACK_TARGET_CHARS);
}

export function detectTokenSmoothingChunk(content: string, locale = "zh") {
  if (!content) {
    return null;
  }

  if (isCjkSmoothingContent(content, locale)) {
    return detectCjkFallbackChunk(content, locale);
  }

  const wordChunk = content.match(/^\S+\s*/u)?.[0];
  if (wordChunk) {
    return wordChunk;
  }

  return sliceCodePoints(content, 1);
}

export function createTokenSmoother({
  enabled,
  delayMs,
  locale,
  maxWaitMs = 100,
  sessionId,
  write,
}: TokenSmootherOptions) {
  const emitChunk = (content: string) => {
    write("token", {
      session_id: sessionId,
      content,
    });
  };

  let buffer = "";
  let bufferedAt = 0;
  let timer: ReturnType<typeof setTimeout> | undefined;
  let scheduledAt = 0;
  let closed = false;
  let releasedChars = 0;
  const barriers: Array<{ target: number; resolve: () => void }> = [];

  function clearTimer() {
    if (timer !== undefined) clearTimeout(timer);
    timer = undefined;
  }

  function release(content: string) {
    buffer = buffer.slice(content.length);
    // The short-tail deadline starts at the last visible output. It does not
    // force a whole paragraph out while regular chunks are still available.
    bufferedAt = buffer ? Date.now() : 0;
    emitChunk(content);
    releasedChars += content.length;
    while (barriers.length && barriers[0].target <= releasedChars) barriers.shift()!.resolve();
  }

  function schedule() {
    if (!buffer) return;
    const now = Date.now();
    const remaining = Math.max(0, maxWaitMs - (now - bufferedAt));
    const wait = detectTokenSmoothingChunk(buffer, locale) || barriers.length ? delayMs : remaining;
    // More input may make a chunk ready sooner, but must never postpone an
    // already scheduled release. A stage boundary uses the same pacing.
    if (timer !== undefined && scheduledAt <= now + wait) return;
    clearTimer();
    scheduledAt = now + wait;
    timer = setTimeout(() => {
      timer = undefined;
      const chunk = detectTokenSmoothingChunk(buffer, locale) ?? buffer;
      const boundary = barriers[0]?.target;
      // Later narration must not be merged across a pending UI/tool boundary.
      release(boundary === undefined ? chunk : chunk.slice(0, boundary - releasedChars));
      schedule();
    }, wait);
  }

  async function flush() {
    if (!buffer) return;
    const target = releasedChars + buffer.length;
    await new Promise<void>((resolve) => {
      barriers.push({ target, resolve });
      schedule();
    });
  }

  return {
    writeToken(content: string) {
      if (!content || closed) return;
      if (!enabled) { emitChunk(content); return; }
      if (!buffer) bufferedAt = Date.now();
      buffer += content;
      schedule();
    },
    // Stage boundaries drain narration without ending future text output.
    flush,
    async close({ immediate = false }: { immediate?: boolean } = {}) {
      closed = true;
      if (immediate) {
        clearTimer();
        if (buffer) release(buffer);
      } else await flush();
    },
  };
}