Jyotisha/frontend/src/lib/session-context-summary.ts

import { Agent } from "@mastra/core/agent";

import {
  consultationHistoryCheckpointChars,
  lastConsultationPair,
  parseSessionContextSummary,
  storedConsultationTurns,
  type SessionContextSummaryV1,
} from "@/lib/consultation-session-history";
import type { ResolvedLanguageModel } from "@/mastra/model";

export const SESSION_CONTEXT_SUMMARY_TIMEOUT_MS = 15_000;
export const SESSION_CONTEXT_SUMMARY_MAX_HAN = 800;
export const SESSION_CONTEXT_SUMMARY_MAX_TOKENS = 600;

export const SESSION_CONTEXT_SUMMARY_INSTRUCTIONS = `你在维护咨询会话的滚动摘要，只供下一轮模型使用。
只根据给定的问答文本写作，不要发明没出现过的事实。
输出不超过 800 个汉字，必须使用下面四个标题，每个标题下用短句：
已问过的问题
已给出的结论（含应期、置信度、blocked 项）
用户补充的事实
未决与待追问
不要写出生日期、出生时间、出生地、姓名、邮箱。`;

type DisposableAbort = Readonly<{
  signal: AbortSignal;
  dispose: () => void;
}>;

function composedAbortSignal(signal: AbortSignal | undefined, timeoutMs: number): DisposableAbort {
  const controller = new AbortController();
  // Must stay ref'd. The platform timeout helper uses an unref timer (BUG-523).
  const timeoutId = globalThis.setTimeout(() => {
    if (!controller.signal.aborted) {
      controller.abort(new DOMException("session context summary timed out", "TimeoutError"));
    }
  }, timeoutMs);
  const onExternalAbort = () => {
    if (!controller.signal.aborted) {
      controller.abort(signal?.reason ?? new DOMException("aborted", "AbortError"));
    }
  };
  if (signal) {
    if (signal.aborted) onExternalAbort();
    else signal.addEventListener("abort", onExternalAbort);
  }
  return {
    signal: controller.signal,
    dispose: () => {
      globalThis.clearTimeout(timeoutId);
      signal?.removeEventListener("abort", onExternalAbort);
    },
  };
}

function whenAborted(signal: AbortSignal): { promise: Promise<never>; dispose: () => void } {
  let onAbort: (() => void) | undefined;
  const promise = new Promise<never>((_, reject) => {
    const fail = () => {
      reject(signal.reason ?? new Error("aborted"));
    };
    if (signal.aborted) {
      fail();
      return;
    }
    onAbort = fail;
    signal.addEventListener("abort", fail, { once: true });
  });
  return {
    promise,
    dispose: () => {
      if (onAbort) signal.removeEventListener("abort", onAbort);
    },
  };
}

function clipHan(value: string, maxChars: number): string {
  const characters = Array.from(value);
  return characters.length > maxChars ? characters.slice(0, maxChars).join("") : value;
}

const ISO_DATE = /\d{4}-\d{2}-\d{2}/g;
const CLOCK_TIME = /\b\d{1,2}:\d{2}(?::\d{2})?\b/g;
const EMAIL = /[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}/gi;

export function sanitizeSessionContextSummary(raw: string): string | null {
  const stripped = raw
    .replace(EMAIL, "")
    .replace(ISO_DATE, "")
    .replace(CLOCK_TIME, "")
    .replace(/[ \t]+\n/g, "\n")
    .replace(/\n{3,}/g, "\n\n")
    .trim();
  if (!stripped) return null;
  return clipHan(stripped, SESSION_CONTEXT_SUMMARY_MAX_HAN);
}

export function tailCharCount(
  messages: unknown,
  summary: SessionContextSummaryV1 | null,
  options: { excludeRequestId?: string; contextWindow?: number | null } = {},
): number {
  const turns = storedConsultationTurns(messages, { excludeRequestId: options.excludeRequestId });
  const tail = summary
    ? turns.filter((turn) => turn.index > summary.throughMessageIndex)
    : turns;
  return tail.reduce((sum, turn) => sum + turn.text.length, 0);
}

export function shouldCheckpoint(
  messages: unknown,
  summary: SessionContextSummaryV1 | null,
  options: { excludeRequestId?: string; contextWindow?: number | null } = {},
): boolean {
  return tailCharCount(messages, summary, options)
    > consultationHistoryCheckpointChars(options.contextWindow);
}

export function messagesForSummaryInput(
  messages: unknown,
  summary: SessionContextSummaryV1 | null,
  options: { excludeRequestId?: string } = {},
): Array<{ role: "user" | "assistant"; text: string; index: number; requestId: string | null }> {
  const turns = storedConsultationTurns(messages, { excludeRequestId: options.excludeRequestId });
  const tail = summary
    ? turns.filter((turn) => turn.index > summary.throughMessageIndex)
    : turns;
  const lastPair = lastConsultationPair(tail);
  const pairStart = lastPair[0];
  if (!pairStart) return [];
  return tail.filter((turn) => turn.index < pairStart.index);
}

export function buildSummaryPrompt(
  previous: SessionContextSummaryV1 | null,
  messages: unknown,
  options: { excludeRequestId?: string } = {},
): string {
  const input = messagesForSummaryInput(messages, previous, options);
  const lines = input.map((turn) => `${turn.role === "user" ? "用户" : "助手"}：${turn.text}`);
  return [
    previous?.text.trim() ? `上一份摘要：\n${previous.text.trim()}` : "上一份摘要：无",
    "需要并入摘要的问答（不含最后一对）：",
    lines.join("\n") || "（无）",
  ].join("\n\n");
}

export async function generateSessionContextSummaryText(
  model: ResolvedLanguageModel,
  prompt: string,
  signal?: AbortSignal,
): Promise<string> {
  const agent = new Agent({
    id: `session-context-summary-${model.id}`,
    name: "Session Context Summary",
    model: model.model,
    instructions: SESSION_CONTEXT_SUMMARY_INSTRUCTIONS,
  });
  const result = await agent.generate([{ role: "user", content: prompt }], {
    abortSignal: signal,
    modelSettings: { maxOutputTokens: SESSION_CONTEXT_SUMMARY_MAX_TOKENS },
  });
  return typeof result.text === "string" ? result.text : "";
}

export async function generateSessionContextSummary(input: {
  model?: ResolvedLanguageModel | null;
  previous: SessionContextSummaryV1 | null;
  messages: unknown;
  excludeRequestId?: string;
  signal?: AbortSignal;
  timeoutMs?: number;
  generateText?: (prompt: string, signal?: AbortSignal) => Promise<string>;
}): Promise<string | null> {
  const prompt = buildSummaryPrompt(input.previous, input.messages, {
    excludeRequestId: input.excludeRequestId,
  });
  const generate = input.generateText ?? (input.model
    ? (nextPrompt: string, signal?: AbortSignal) => generateSessionContextSummaryText(
      input.model as ResolvedLanguageModel,
      nextPrompt,
      signal,
    )
    : null);
  if (!generate) return null;
  const composed = composedAbortSignal(input.signal, input.timeoutMs ?? SESSION_CONTEXT_SUMMARY_TIMEOUT_MS);
  const aborted = whenAborted(composed.signal);
  try {
    const raw = await Promise.race([generate(prompt, composed.signal), aborted.promise]);
    return sanitizeSessionContextSummary(raw);
  } catch {
    return null;
  } finally {
    aborted.dispose();
    composed.dispose();
  }
}

export function nextSessionContextSummary(
  messages: unknown,
  previous: SessionContextSummaryV1 | null,
  text: string,
  updatedAt: string,
  options: { excludeRequestId?: string } = {},
): SessionContextSummaryV1 | null {
  const covered = messagesForSummaryInput(messages, previous, options);
  const last = covered.at(-1);
  if (!last) return null;
  return {
    version: 1,
    text,
    throughRequestId: last.requestId ?? previous?.throughRequestId ?? "",
    throughMessageIndex: last.index,
    messageCount: last.index + 1,
    updatedAt,
  };
}

export async function resolveInheritedContextSummary(input: {
  continuedFromSessionId?: string | null;
  loadOwnedSummary: (sessionId: string) => Promise<unknown>;
}): Promise<SessionContextSummaryV1 | null> {
  const sourceId = input.continuedFromSessionId?.trim();
  if (!sourceId) return null;
  try {
    return parseSessionContextSummary(await input.loadOwnedSummary(sourceId));
  } catch {
    return null;
  }
}

export async function writeSessionContextSummary(input: {
  seenUpdatedAt: string | null;
  summary: SessionContextSummaryV1;
  update: (summary: SessionContextSummaryV1, seenUpdatedAt: string | null) => Promise<boolean>;
}): Promise<"written" | "abandoned"> {
  const written = await input.update(input.summary, input.seenUpdatedAt);
  return written ? "written" : "abandoned";
}

export async function checkpointSessionContextSummary(input: {
  messages: unknown;
  summary: unknown;
  excludeRequestId?: string;
  contextWindow?: number | null;
  now?: () => Date;
  generateText: (prompt: string, signal?: AbortSignal) => Promise<string>;
  update: (summary: SessionContextSummaryV1, seenUpdatedAt: string | null) => Promise<boolean>;
  timeoutMs?: number;
}): Promise<"written" | "skipped" | "abandoned" | "failed"> {
  const previous = parseSessionContextSummary(input.summary);
  if (!shouldCheckpoint(input.messages, previous, {
    excludeRequestId: input.excludeRequestId,
    contextWindow: input.contextWindow,
  })) {
    return "skipped";
  }
  try {
    const text = await generateSessionContextSummary({
      previous,
      messages: input.messages,
      excludeRequestId: input.excludeRequestId,
      generateText: input.generateText,
      timeoutMs: input.timeoutMs,
    });
    if (!text) return "failed";
    const next = nextSessionContextSummary(
      input.messages,
      previous,
      text,
      (input.now ?? (() => new Date()))().toISOString(),
      { excludeRequestId: input.excludeRequestId },
    );
    if (!next) return "skipped";
    return writeSessionContextSummary({
      seenUpdatedAt: previous?.updatedAt ?? null,
      summary: next,
      update: input.update,
    });
  } catch {
    return "failed";
  }
}