Jyotisha/frontend/tests/consult-natal-agent-test-support.ts

// Shared harness for natal-route tests that drive the real personal Agent
// (`getJyotishAgent`, real skill binding, real calculation and lookup tools,
// real prepareStep, real run clock) over a prompt-recording fake model.
// Same wiring as consult-single-pass-answer-20260927.test.ts (BUG-1053); the
// workflow the calculation tool receives is a golden engine capture.
import { getJyotishAgent } from "../src/mastra/index.ts";
import { createNdjsonParser } from "../src/lib/consultation-agent-events.ts";
import {
  consultationContinueMessages,
  natalAnswerShapeInstruction,
} from "../src/lib/consultation-thinking-plan.ts";
import { streamAgentResponse } from "../src/lib/stream-agent-response.ts";
import type { ConsultationDomain } from "../src/lib/consultation-domain-registry.ts";
import {
  AGENT_MAX_STEPS,
  AGENT_TIMEOUT_MS,
  CONSULTATION_ANSWER_TIMEOUT_MS,
  consultationNatalPrepareStep,
  consultationStepBudgetReceipt,
  createConsultationAgentContext,
  createConsultationRunClock,
  createConsultationRuntimeState,
  publicConsultationRuntimeSteps,
} from "../src/mastra/consultation-tools.ts";

export type Part = { text?: string; tool?: { name: string; input: Record<string, unknown> } };
export type Turn = { parts: Part[]; finish: string; delayMs?: number };

/** A LanguageModelV2 that plays `turns` in order and records every prompt. */
export function scriptedModel(turns: Turn[]) {
  const prompts: unknown[][] = [];
  let call = 0;
  const model = {
    specificationVersion: "v2",
    provider: "fake",
    modelId: "fake-natal-agent",
    supportedUrls: {},
    async doGenerate() {
      throw new Error("not used");
    },
    async doStream(options: { prompt: unknown[]; abortSignal?: AbortSignal }) {
      prompts.push(options.prompt);
      const turn = turns[call] ?? { parts: [], finish: "stop" };
      call += 1;
      const signal = options.abortSignal;
      const stream = new ReadableStream({
        async start(controller) {
          controller.enqueue({ type: "stream-start", warnings: [] });
          let textOpen = false;
          for (const [index, part] of turn.parts.entries()) {
            if (turn.delayMs) {
              try {
                await new Promise<void>((resolve, reject) => {
                  if (signal?.aborted) return reject(signal.reason);
                  const timer = setTimeout(resolve, turn.delayMs);
                  signal?.addEventListener("abort", () => {
                    clearTimeout(timer);
                    reject(signal.reason);
                  }, { once: true });
                });
              } catch (error) {
                controller.error(error);
                return;
              }
            }
            if (part.text !== undefined) {
              if (!textOpen) {
                controller.enqueue({ type: "text-start", id: `t${call}` });
                textOpen = true;
              }
              controller.enqueue({ type: "text-delta", id: `t${call}`, delta: part.text });
            }
            if (part.tool) {
              if (textOpen) {
                controller.enqueue({ type: "text-end", id: `t${call}` });
                textOpen = false;
              }
              controller.enqueue({
                type: "tool-call",
                toolCallId: `call-${call}-${index}`,
                toolName: part.tool.name,
                input: JSON.stringify(part.tool.input),
              });
            }
          }
          if (textOpen) controller.enqueue({ type: "text-end", id: `t${call}` });
          controller.enqueue({
            type: "finish",
            finishReason: turn.finish,
            usage: { inputTokens: 10, outputTokens: 10, totalTokens: 20 },
          });
          controller.close();
        },
      });
      return { stream };
    },
  };
  return { model, prompts, calls: () => call };
}

export function pieces(text: string, size = 12) {
  return (text.match(new RegExp(`[\\s\\S]{1,${size}}`, "g")) ?? []).map((value) => ({ text: value }));
}

export const publicServerChart = {
  name: "public",
  toolInput: {
    year: 1955, month: 2, day: 24, hour: 19, minute: 15, city: "San Francisco", lat: 37.77, lon: -122.42, tz: -8,
    ayanamsa: "raman" as const, declared_accuracy: "minute" as const, time_source: "aa_rated",
  },
  truth: {
    birthDate: "1955-02-24", reportedBirthTime: "19:15", activeBirthTime: null,
    selectedTimeKind: "reported" as const, birthTimeSource: "reported", birthTimeStatus: "reported",
    placeLabel: "San Francisco", placeCodes: { countryCode: "US", provinceCode: null, cityCode: null, districtCode: null },
    placeId: null, placeType: "city", placeProvider: "profile", latitude: 37.77, longitude: -122.42,
    timezoneId: "America/Los_Angeles", timezoneSource: "profile", timezoneOffset: -8,
  },
};

let seq = 0;

/**
 * The natal route's wiring, minus HTTP, auth and billing: the same Agent,
 * prepareStep, run clock, stream options, step-scoped answer, answer-phase
 * hand-over, continuation builder and answer retry as `runAgenticConsultation`.
 */
export async function runNatalAgent(
  turns: Turn[],
  options: {
    workflow: Record<string, unknown>;
    theme?: ConsultationDomain;
    question?: string;
    toolPhaseMs?: number;
    answerMs?: number;
  },
) {
  seq += 1;
  const { model, prompts, calls } = scriptedModel(turns);
  const state = createConsultationRuntimeState({ plannedSteps: AGENT_MAX_STEPS });
  const clock = createConsultationRunClock({
    toolPhaseMs: options.toolPhaseMs ?? AGENT_TIMEOUT_MS,
    answerMs: options.answerMs ?? CONSULTATION_ANSWER_TIMEOUT_MS,
    answerReady: () => state.consultationToolCompleted,
  });
  let workflowRuns = 0;
  const answerSignals: AbortSignal[] = [];
  const agentContext = createConsultationAgentContext({
    userId: "u", sessionId: "s", requestId: `natal-agent-${seq}`, consultationMode: "verified_chart",
    theme: options.theme ?? "parents", serverChart: publicServerChart as never, abortSignal: clock.toolSignal, state,
    runWorkflow: async () => {
      workflowRuns += 1;
      return structuredClone(options.workflow) as never;
    },
  });
  const agent = getJyotishAgent({ id: `fake-${seq}`, model } as never, agentContext);
  const question = options.question ?? "我和父母关系如何";
  const baseMessages = [{ role: "user" as const, content: `${natalAnswerShapeInstruction()}\n问题：${question}` }];
  const streamOptions = { runId: `run-${seq}`, maxSteps: AGENT_MAX_STEPS, abortSignal: clock.loopSignal };
  const natalStreamOptions = { ...streamOptions, prepareStep: consultationNatalPrepareStep };
  const continuations: unknown[] = [];
  let completed: string | null = null;
  let charges = 0;
  let errored: unknown = null;
  const result = await agent.stream(baseMessages as never, natalStreamOptions as never);
  const response = streamAgentResponse({
    runId: `run-${seq}`,
    requestId: `req-${seq}`,
    state,
    stream: result.fullStream as ReadableStream<unknown>,
    requireTool: true,
    stepScopedAnswer: true,
    onAnswerPhase: () => { answerSignals.push(clock.answerSignal()); },
    pass4Mode: "verified_chart",
    retryForAnswer: async (retryHint) => {
      const retried = await agent.stream([
        ...baseMessages,
        { role: "user" as const, content: `服务器计算已经完成，但上一轮没有输出任何回答文本。请重新取回本次计算结果，然后直接给出回答。${retryHint ? `\n${retryHint}` : ""}` },
      ] as never, { ...natalStreamOptions, abortSignal: clock.answerSignal() } as never);
      return retried.fullStream as ReadableStream<unknown>;
    },
    continueAfterLength: async (output, evidence) => {
      continuations.push(evidence);
      answerSignals.push(clock.answerSignal());
      const continued = await agent.stream(
        consultationContinueMessages(baseMessages, output, evidence) as never,
        { ...streamOptions, abortSignal: clock.answerSignal() } as never,
      );
      return continued.fullStream as ReadableStream<unknown>;
    },
    toolStatus: () => "ready",
    receipt: () => ({
      runId: `run-${seq}`,
      runtime: "mastra-agentic",
      skill: { name: "jyotish-vedic-astrology", loaded: true, referenceReads: 0, methodologySections: 0 },
      steps: publicConsultationRuntimeSteps(state),
      stepBudget: consultationStepBudgetReceipt(state),
      workflow: state.workflowReceipt ?? { route: "pending", status: "blocked", preciseTiming: "blocked", missingLayers: [] },
      techniqueTruth: "unknown",
    }) as never,
    onComplete: (output) => { completed = output; charges += 1; },
    onError: (error) => { errored = error; },
  });
  const events: Array<{ type: string; code?: string; text?: string; label?: string; phase?: string; receipt?: unknown }> = [];
  const parser = createNdjsonParser((event) => events.push(event as never));
  parser.finish(await response.text());
  const answer = events.filter((event) => event.type === "answer.delta").map((event) => event.text ?? "").join("");
  const terminal = events.filter((event) => event.type === "run.completed" || event.type === "run.failed");
  return {
    state, events, answer, terminal, prompts, calls: calls(), continuations, workflowRuns, clock, answerSignals,
    completed: completed as string | null, charges, errored,
  };
}

/** Index of the first prompt that already contains a tool result for `toolName`. */
export function firstPromptWithToolResult(prompts: unknown[][], toolName: string) {
  return prompts.findIndex((prompt) => prompt.some((message) => {
    const value = message as { role?: string; content?: unknown };
    return value.role === "tool" && JSON.stringify(value.content).includes(toolName);
  }));
}