nuxt-app/server/service/agent/stream-buffer.ts

import log4js from "logger";

const logger = log4js.getLogger("APP");

const BUFFER_TTL_MS = 5 * 60 * 1000;
const CLEANUP_INTERVAL_MS = 60 * 1000;

export interface StreamBuffer {
  chunks: Uint8Array[];
  done: boolean;
  createdAt: number;
  doneAt: number | null;
  modelId: number | null;
  userMessageId: string | null;
  subscribers: Array<(chunk: Uint8Array) => void>;
  doneSubscribers: Array<() => void>;
}

const buffers = new Map<string, StreamBuffer>();

let cleanupTimer: ReturnType<typeof setInterval> | null = null;

function ensureCleanupTimer() {
  if (cleanupTimer) return;
  cleanupTimer = setInterval(() => {
    const now = Date.now();
    for (const [sid, buf] of buffers) {
      const expireAt = buf.doneAt ?? buf.createdAt + BUFFER_TTL_MS;
      if (now > expireAt) {
        buffers.delete(sid);
        logger.info("[STREAM-BUFFER] expired and removed sessionId=%s", sid);
      }
    }
    if (buffers.size === 0 && cleanupTimer) {
      clearInterval(cleanupTimer);
      cleanupTimer = null;
    }
  }, CLEANUP_INTERVAL_MS);
  cleanupTimer.unref?.();
}

export function createStreamBuffer(
  sessionId: string,
  meta: { modelId?: number | null; userMessageId?: string | null },
): StreamBuffer {
  const existing = buffers.get(sessionId);
  if (existing && !existing.done) {
    existing.modelId = meta.modelId ?? existing.modelId;
    existing.userMessageId = meta.userMessageId ?? existing.userMessageId;
    logger.info("[STREAM-BUFFER] reused existing sessionId=%s chunks=%d", sessionId, existing.chunks.length);
    return existing;
  }
  const buf: StreamBuffer = {
    chunks: [],
    done: false,
    createdAt: Date.now(),
    doneAt: null,
    modelId: meta.modelId ?? null,
    userMessageId: meta.userMessageId ?? null,
    subscribers: [],
    doneSubscribers: [],
  };
  buffers.set(sessionId, buf);
  ensureCleanupTimer();
  logger.info("[STREAM-BUFFER] created sessionId=%s", sessionId);
  return buf;
}

export function appendChunk(sessionId: string, chunk: Uint8Array): void {
  const buf = buffers.get(sessionId);
  if (!buf || buf.done) return;
  buf.chunks.push(chunk);
  for (const sub of buf.subscribers) {
    try {
      sub(chunk);
    } catch (e) {
      logger.error("[STREAM-BUFFER] subscriber error: %s", e instanceof Error ? e.message : String(e));
    }
  }
}

export function markBufferDone(sessionId: string): void {
  const buf = buffers.get(sessionId);
  if (!buf) return;
  buf.done = true;
  buf.doneAt = Date.now();
  const doneSubs = buf.doneSubscribers;
  buf.doneSubscribers = [];
  buf.subscribers = [];
  for (const sub of doneSubs) {
    try {
      sub();
    } catch (e) {
      logger.error("[STREAM-BUFFER] done subscriber error: %s", e instanceof Error ? e.message : String(e));
    }
  }
  logger.info("[STREAM-BUFFER] done sessionId=%s chunks=%d", sessionId, buf.chunks.length);
}

export function getStreamBuffer(sessionId: string): StreamBuffer | undefined {
  return buffers.get(sessionId);
}

export function removeStreamBuffer(sessionId: string): void {
  buffers.delete(sessionId);
}

export function hasActiveStream(sessionId: string): boolean {
  const buf = buffers.get(sessionId);
  return !!buf && !buf.done;
}

export function subscribeToBuffer(
  sessionId: string,
  onChunk: (chunk: Uint8Array) => void,
  onDone: () => void,
): () => void {
  const buf = buffers.get(sessionId);
  if (!buf) {
    onDone();
    return () => {};
  }

  if (buf.done) {
    onDone();
    return () => {};
  }

  buf.subscribers.push(onChunk);
  buf.doneSubscribers.push(onDone);

  return () => {
    buf.subscribers = buf.subscribers.filter((s) => s !== onChunk);
    buf.doneSubscribers = buf.doneSubscribers.filter((s) => s !== onDone);
  };
}