ai-chatbot/lib/agents/psd-remote-asset.ts

import "server-only";

const DEFAULT_TIMEOUT_MS = 30_000;
const DEFAULT_RETRY_COUNT = 3;
const RETRY_DELAY_MS = 1_000;
const DEFAULT_MEMORY_CACHE_TTL_MS = 10 * 60 * 1000;
const DEFAULT_MEMORY_CACHE_MAX_BYTES = 256 * 1024 * 1024;
const DEFAULT_MEMORY_CACHE_MAX_ENTRY_BYTES = 32 * 1024 * 1024;
const remoteAssetCacheKey = "__agent_psd_remote_asset_cache__";

type RemoteAssetCacheEntry = {
  buffer: Buffer;
  byteLength: number;
  expiresAt: number;
  lastAccessAt: number;
};

type RemoteAssetCacheState = {
  entries: Map<string, RemoteAssetCacheEntry>;
  inFlight: Map<string, Promise<Buffer>>;
  totalBytes: number;
};

function getRemoteAssetCacheState(): RemoteAssetCacheState {
  const globalObj = globalThis as typeof globalThis & {
    [remoteAssetCacheKey]?: RemoteAssetCacheState;
  };

  if (!globalObj[remoteAssetCacheKey]) {
    globalObj[remoteAssetCacheKey] = {
      entries: new Map<string, RemoteAssetCacheEntry>(),
      inFlight: new Map<string, Promise<Buffer>>(),
      totalBytes: 0,
    };
  }

  return globalObj[remoteAssetCacheKey] as RemoteAssetCacheState;
}

function parsePositiveInteger(value: string | undefined, fallback: number) {
  const parsed = Number.parseInt(String(value || ""), 10);
  return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback;
}

function getRemoteAssetCacheConfig() {
  return {
    ttlMs: parsePositiveInteger(
      process.env.AGENT_PSD_REMOTE_ASSET_CACHE_TTL_MS,
      DEFAULT_MEMORY_CACHE_TTL_MS
    ),
    maxBytes: parsePositiveInteger(
      process.env.AGENT_PSD_REMOTE_ASSET_CACHE_MAX_BYTES,
      DEFAULT_MEMORY_CACHE_MAX_BYTES
    ),
    maxEntryBytes: parsePositiveInteger(
      process.env.AGENT_PSD_REMOTE_ASSET_CACHE_MAX_ENTRY_BYTES,
      DEFAULT_MEMORY_CACHE_MAX_ENTRY_BYTES
    ),
  };
}

function isRemoteAssetMemoryCacheEnabled() {
  return process.env.AGENT_PSD_REMOTE_ASSET_CACHE_DISABLED !== "1";
}

function isCacheableRemoteAssetUrl(url: string) {
  return !/^(data|blob):/i.test(String(url || "").trim());
}

function pruneRemoteAssetCache(state: RemoteAssetCacheState, maxBytes: number) {
  const now = Date.now();
  for (const [key, entry] of state.entries) {
    if (entry.expiresAt <= now) {
      state.entries.delete(key);
      state.totalBytes -= entry.byteLength;
    }
  }

  if (state.totalBytes <= maxBytes) {
    return;
  }

  const entriesByLastAccess = Array.from(state.entries.entries()).sort(
    (left, right) => left[1].lastAccessAt - right[1].lastAccessAt
  );

  for (const [key, entry] of entriesByLastAccess) {
    if (state.totalBytes <= maxBytes) break;
    state.entries.delete(key);
    state.totalBytes -= entry.byteLength;
  }
}

function readRemoteAssetMemoryCache(url: string) {
  if (!isRemoteAssetMemoryCacheEnabled() || !isCacheableRemoteAssetUrl(url)) {
    return null;
  }

  const state = getRemoteAssetCacheState();
  const entry = state.entries.get(url);
  if (!entry) {
    return null;
  }

  if (entry.expiresAt <= Date.now()) {
    state.entries.delete(url);
    state.totalBytes -= entry.byteLength;
    return null;
  }

  entry.lastAccessAt = Date.now();
  return entry.buffer;
}

function writeRemoteAssetMemoryCache(url: string, buffer: Buffer) {
  if (!isRemoteAssetMemoryCacheEnabled() || !isCacheableRemoteAssetUrl(url)) {
    return;
  }

  const { ttlMs, maxBytes, maxEntryBytes } = getRemoteAssetCacheConfig();
  const byteLength = buffer.byteLength;
  if (byteLength <= 0 || byteLength > maxEntryBytes) {
    return;
  }

  const state = getRemoteAssetCacheState();
  const previous = state.entries.get(url);
  if (previous) {
    state.totalBytes -= previous.byteLength;
  }

  state.entries.set(url, {
    buffer,
    byteLength,
    expiresAt: Date.now() + ttlMs,
    lastAccessAt: Date.now(),
  });
  state.totalBytes += byteLength;
  pruneRemoteAssetCache(state, maxBytes);
}

function sleep(ms: number) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

export function describeRemoteAssetLocation(url: string) {
  const normalizedUrl = String(url || "").trim();
  if (!normalizedUrl) {
    return "未知地址";
  }

  try {
    const parsed = new URL(normalizedUrl);
    return parsed.host || parsed.hostname || normalizedUrl;
  } catch {
    return normalizedUrl.length > 80 ? `${normalizedUrl.slice(0, 77)}...` : normalizedUrl;
  }
}

function isRetryableStatus(status: number) {
  return status === 408 || status === 425 || status === 429 || status >= 500;
}

function isRetryableFetchError(error: unknown) {
  if (!(error instanceof Error)) {
    return true;
  }

  const message = `${error.name} ${error.message}`.toLowerCase();
  return (
    error.name === "AbortError" ||
    message.includes("fetch failed") ||
    message.includes("timeout") ||
    message.includes("timed out") ||
    message.includes("econnreset") ||
    message.includes("socket") ||
    message.includes("network") ||
    message.includes("terminated") ||
    message.includes("und_err")
  );
}

export async function fetchRemoteAssetBuffer(
  url: string,
  options?: {
    label?: string;
    timeoutMs?: number;
    retries?: number;
    useMemoryCache?: boolean;
  }
) {
  const normalizedUrl = String(url || "").trim();
  const label = String(options?.label || "远程资源").trim() || "远程资源";
  if (!normalizedUrl) {
    throw new Error(`${label}地址为空`);
  }

  const attempts = Math.max(1, Math.round(Number(options?.retries || DEFAULT_RETRY_COUNT)));
  const timeoutMs = Math.max(5_000, Math.round(Number(options?.timeoutMs || DEFAULT_TIMEOUT_MS)));
  const location = describeRemoteAssetLocation(normalizedUrl);
  const shouldUseMemoryCache = options?.useMemoryCache !== false;

  if (shouldUseMemoryCache) {
    const cachedBuffer = readRemoteAssetMemoryCache(normalizedUrl);
    if (cachedBuffer) {
      return cachedBuffer;
    }

    const cacheState = getRemoteAssetCacheState();
    const inFlightRequest = cacheState.inFlight.get(normalizedUrl);
    if (inFlightRequest) {
      return inFlightRequest;
    }
  }

  const downloadPromise = (async () => {
    for (let attempt = 1; attempt <= attempts; attempt += 1) {
      try {
        const response = await fetch(normalizedUrl, {
          cache: "no-store",
          signal: AbortSignal.timeout(timeoutMs),
        });

        if (!response.ok) {
          const statusError = new Error(`${label}下载失败（${response.status}）`);
          if (attempt < attempts && isRetryableStatus(response.status)) {
            console.warn("[PSD Remote Asset] transient status, retrying download:", {
              label,
              location,
              status: response.status,
              attempt,
              attempts,
            });
            await sleep(RETRY_DELAY_MS * attempt);
            continue;
          }

          throw statusError;
        }

        const buffer = Buffer.from(await response.arrayBuffer());
        if (shouldUseMemoryCache) {
          writeRemoteAssetMemoryCache(normalizedUrl, buffer);
        }
        return buffer;
      } catch (error) {
        if (attempt < attempts && isRetryableFetchError(error)) {
          console.warn("[PSD Remote Asset] transient fetch error, retrying download:", {
            label,
            location,
            attempt,
            attempts,
            error,
          });
          await sleep(RETRY_DELAY_MS * attempt);
          continue;
        }

        console.error("[PSD Remote Asset] failed to download asset:", {
          label,
          location,
          attemptsTried: attempt,
          error,
        });
        throw new Error(`${label}下载失败，请稍后重试或检查资源是否可访问（${location}）`);
      }
    }

    throw new Error(`${label}下载失败，请稍后重试或检查资源是否可访问（${location}）`);
  })();

  if (shouldUseMemoryCache && isCacheableRemoteAssetUrl(normalizedUrl)) {
    const cacheState = getRemoteAssetCacheState();
    cacheState.inFlight.set(normalizedUrl, downloadPromise);
    downloadPromise.then(
      () => {
        cacheState.inFlight.delete(normalizedUrl);
      },
      () => {
        cacheState.inFlight.delete(normalizedUrl);
      }
    );
  }

  return downloadPromise;
}