@assistant-ui/react/src/unstable/useMessageStallDetection.ts

Open-source TypeScript/React library for building production-grade AI chat experiences

"use client";

import { useEffect, useRef, useState } from "react";
import { useAuiState } from "@assistant-ui/store";
import { useShallowSelector } from "@assistant-ui/store/internal";

/** @deprecated Experimental since 2026-06-13. Not scheduled for removal; the API may change in any release. */
export type Unstable_MessageStallDetectionOptions = {
  /**
   * Milliseconds of unchanged message content before the message counts as
   * stalled.
   * @default 2000
   */
  thresholdMs?: number | undefined;
};

/** @deprecated Experimental since 2026-06-13. Not scheduled for removal; the API may change in any release. */
export type Unstable_MessageStallDetection = {
  /** True while the message is running and its content has not changed for at least `thresholdMs`. */
  stalled: boolean;
  /** Milliseconds since the last observed content change. `0` while not stalled. */
  stalledForMs: number;
};

/**
 * Detects mid-run output stalls on the current message: while the message is
 * running, watches its text, reasoning, and tool-argument values plus tool-result
 * availability and reports a stall once they stop changing for `thresholdMs`.
 * Useful for re-surfacing a "still working" indicator during tool think-time or
 * provider stalls, after the first tokens have already streamed.
 *
 * Must be used inside a message scope.
 *
 * @deprecated Experimental since 2026-06-13. Not scheduled for removal; the API may change in any release.
 */
export function unstable_useMessageStallDetection(
  options?: Unstable_MessageStallDetectionOptions,
): Unstable_MessageStallDetection {
  const thresholdMs = options?.thresholdMs ?? 2000;

  const activity = useAuiState(
    useShallowSelector((s) => {
      const running = s.message.status?.type === "running";
      if (!running) return [false];

      const values: unknown[] = [true, s.message.content.length];

      for (const part of s.message.content) {
        if (part.type === "text" || part.type === "reasoning") {
          values.push(part.type, part.text);
        } else if (part.type === "tool-call") {
          values.push(part.type, part.argsText, part.result !== undefined);
        } else {
          values.push(part.type);
        }
      }
      return values;
    }),
  );

  const running = activity[0] === true;
  const lastActivityRef = useRef(0);
  const [stalled, setStalled] = useState(false);
  const [now, setNow] = useState(0);

  useEffect(() => {
    if (!running) return undefined;
    lastActivityRef.current = Date.now();
    return undefined;
  }, [running, activity]);

  useEffect(() => {
    if (!running) {
      setStalled(false);
      return undefined;
    }

    const stall = () => {
      setNow(Date.now());
      setStalled(true);
    };
    const sinceActivity = Date.now() - lastActivityRef.current;
    if (sinceActivity >= thresholdMs) {
      stall();
      return undefined;
    }

    setStalled(false);
    const id = setTimeout(stall, thresholdMs - sinceActivity);
    return () => clearTimeout(id);
  }, [running, activity, thresholdMs]);

  useEffect(() => {
    if (!stalled) return undefined;
    const id = setInterval(() => setNow(Date.now()), 1000);
    return () => clearInterval(id);
  }, [stalled]);

  if (!stalled) return { stalled: false, stalledForMs: 0 };
  return {
    stalled: true,
    stalledForMs: Math.max(0, now - lastActivityRef.current),
  };
}