Jyotisha/frontend/src/lib/session-title-agent.ts

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

import { consultationDomainDefinition, type ConsultationDomain } from "@/lib/consultation-domain-registry";
import {
  sanitizeSessionTitle,
  shouldGenerateSessionTitle,
  isAutoDerivedSessionTitle,
} from "@/lib/session-title";
import type { ResolvedLanguageModel } from "@/mastra/model";

export {
  sanitizeSessionTitle,
  shouldGenerateSessionTitle,
  isAutoDerivedSessionTitle,
};

export const SESSION_TITLE_TIMEOUT_MS = 8_000;
export const SESSION_TITLE_MAX_TOKENS = 30;

export const SESSION_TITLE_INSTRUCTIONS = `你只给这段咨询起一个中文标题。
6 到 12 个字，直接写主题，不要引号、书名号、句号或解释。
不要复述出生日期、时间或地点，也不要写资料姓名。
只输出标题本身。`;

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 title 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);
    },
  };
}

export async function generateSessionTitleText(
  model: ResolvedLanguageModel,
  input: {
    question: string;
    theme?: ConsultationDomain | null;
    chartRole?: "self" | "other" | null;
  },
  signal?: AbortSignal,
): Promise<string> {
  const agent = new Agent({
    id: `session-title-${model.id}`,
    name: "Session Title",
    model: model.model,
    instructions: SESSION_TITLE_INSTRUCTIONS,
  });
  const themeLabel = input.theme && input.theme !== "general"
    ? consultationDomainDefinition(input.theme).label
    : "综合";
  const subject = input.chartRole === "other" ? "他人" : "本人";
  const result = await agent.generate([{
    role: "user",
    content: [
      `主题：${themeLabel}`,
      `分析对象：${subject}`,
      `问题：${input.question.trim()}`,
    ].join("\n"),
  }], {
    abortSignal: signal,
    modelSettings: { maxOutputTokens: SESSION_TITLE_MAX_TOKENS },
  });
  const text = typeof result.text === "string" ? result.text : "";
  return text.trim();
}

export async function generateSessionTitle(input: {
  model?: ResolvedLanguageModel | null;
  question: string;
  theme?: ConsultationDomain | null;
  chartRole?: "self" | "other" | null;
  signal?: AbortSignal;
  timeoutMs?: number;
  generateText?: (signal?: AbortSignal) => Promise<string>;
}): Promise<string | null> {
  const generate = input.generateText ?? (input.model
    ? (signal?: AbortSignal) => generateSessionTitleText(input.model as ResolvedLanguageModel, {
      question: input.question,
      theme: input.theme,
      chartRole: input.chartRole,
    }, signal)
    : null);
  if (!generate) return null;
  const composed = composedAbortSignal(input.signal, input.timeoutMs ?? SESSION_TITLE_TIMEOUT_MS);
  const aborted = whenAborted(composed.signal);
  try {
    const raw = await Promise.race([generate(composed.signal), aborted.promise]);
    return sanitizeSessionTitle(raw);
  } catch {
    return null;
  } finally {
    aborted.dispose();
    composed.dispose();
  }
}