Jyotisha/frontend/tests/answer-effort-probe.mts

/**
 * Request-body probe for issue #65.
 *
 * --dry-run needs no model key. It sends one agent.stream per shape through
 * Mastra's OpenAI-compatible channel and prints the thinking / reasoning_effort
 * / max_tokens fields of each request. A fake fetch never leaves the machine.
 *
 * Shapes:
 *   (a) consultationGenerationSettings (today's answer settings)
 *   (b) consultationSliceGenerationSettings (nested reasoningEffort, no caller)
 *   (c) providerReasoningEffort low / medium / high / max at the top level
 *
 * With DS_KEY or BACKTEST_MODEL_KEY, the same file can ask the provider which
 * efforts it accepts. This checkout does not run that half.
 */
import { writeFileSync } from "node:fs";

import { Agent } from "@mastra/core/agent";
import { createTool } from "@mastra/core/tools";
import { z } from "zod";

import { agentGenerationSettings } from "../src/lib/agent-generation-settings.ts";
import {
  consultationContinueGenerationSettings,
  consultationGenerationSettings,
  consultationNatalPrepareStep,
  consultationSliceGenerationSettings,
} from "../src/mastra/consultation-tools.ts";

const MODEL = {
  providerId: "deepseek",
  modelId: "deepseek-flash",
  url: "http://127.0.0.1:9/v1",
  apiKey: "dummy",
};

const markTool = createTool({
  id: "run-jyotish-consultation",
  description: "Mark the calculation done.",
  inputSchema: z.object({}).strict(),
  execute: async () => ({ ok: true }),
});

function sse(chunks: unknown[]) {
  const body = `${chunks.map((chunk) => `data: ${JSON.stringify(chunk)}\n\n`).join("")}data: [DONE]\n\n`;
  return new Response(body, {
    status: 200,
    headers: { "content-type": "text/event-stream" },
  });
}

function toolCallChunk() {
  return sse([
    {
      id: "chatcmpl-tool",
      object: "chat.completion.chunk",
      choices: [{
        index: 0,
        delta: {
          role: "assistant",
          tool_calls: [{
            index: 0,
            id: "call_1",
            type: "function",
            function: { name: "run-jyotish-consultation", arguments: "" },
          }],
        },
        finish_reason: null,
      }],
    },
    {
      id: "chatcmpl-tool",
      object: "chat.completion.chunk",
      choices: [{
        index: 0,
        delta: { tool_calls: [{ index: 0, function: { arguments: "{}" } }] },
        finish_reason: null,
      }],
    },
    {
      id: "chatcmpl-tool",
      object: "chat.completion.chunk",
      choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }],
    },
  ]);
}

function textChunk() {
  return sse([
    {
      id: "chatcmpl-text",
      object: "chat.completion.chunk",
      choices: [{ index: 0, delta: { role: "assistant", content: "ok" }, finish_reason: null }],
    },
    {
      id: "chatcmpl-text",
      object: "chat.completion.chunk",
      choices: [{ index: 0, delta: {}, finish_reason: "stop" }],
      usage: { prompt_tokens: 1, completion_tokens: 1, completion_tokens_details: { reasoning_tokens: 0 } },
    },
  ]);
}

export function requestFields(body: unknown) {
  const record = body && typeof body === "object" ? body as Record<string, unknown> : {};
  return {
    thinking: record.thinking ?? null,
    reasoning_effort: record.reasoning_effort ?? null,
    max_tokens: record.max_tokens ?? null,
  };
}

export async function captureStreamBodies(input: {
  settings: ReturnType<typeof consultationGenerationSettings>;
  prepareStep?: (step: { stepNumber: number }) => { toolChoice?: "auto"; activeTools?: string[]; providerOptions?: unknown };
  steps: 1 | 2;
}): Promise<unknown[]> {
  const bodies: unknown[] = [];
  const previous = globalThis.fetch;
  globalThis.fetch = async (_url, init) => {
    const raw = init?.body;
    bodies.push(JSON.parse(typeof raw === "string" ? raw : "{}"));
    return bodies.length === 1 && input.steps === 2 ? toolCallChunk() : textChunk();
  };
  try {
    const agent = new Agent({
      id: `answer-effort-probe-${bodies.length}-${Date.now()}`,
      name: "answer-effort-probe",
      instructions: "Reply ok.",
      model: MODEL,
      tools: { "run-jyotish-consultation": markTool },
    });
    const result = await agent.stream([{ role: "user", content: "ping" }], {
      maxSteps: input.steps,
      ...input.settings,
      ...(input.prepareStep ? { prepareStep: input.prepareStep } : {}),
    } as never);
    const drained: unknown[] = [];
    for await (const chunk of result.fullStream) drained.push(chunk);
    if (drained.length === 0) throw new Error("probe stream produced no chunks");
  } finally {
    globalThis.fetch = previous;
  }
  return bodies;
}

async function main() {
  if (!process.argv.includes("--dry-run")) {
    console.log("A/B 未跑：缺模型 key。只做请求体拦截时加上 --dry-run。");
    return;
  }
  const model = MODEL;
  const natal = await captureStreamBodies({
    settings: consultationGenerationSettings(model),
    prepareStep: (input) => consultationNatalPrepareStep({
      stepNumber: input.stepNumber,
      answerReady: input.stepNumber >= 1,
    }),
    steps: 2,
  });
  const slice = await captureStreamBodies({
    settings: consultationSliceGenerationSettings(model),
    steps: 1,
  });
  const topLevel = [];
  for (const effort of ["low", "medium", "high", "max"] as const) {
    const bodies = await captureStreamBodies({
      settings: agentGenerationSettings(model, { thinking: "enabled", providerReasoningEffort: effort }),
      steps: 1,
    });
    topLevel.push({ effort, fields: requestFields(bodies[0]) });
  }
  const continued = await captureStreamBodies({
    settings: consultationContinueGenerationSettings(model),
    steps: 1,
  });
  const report = {
    fixtureProvenance: "frontend/tests/answer-effort-probe.mts --dry-run",
    a_consultationGenerationSettings: natal.map(requestFields),
    b_consultationSliceGenerationSettings: slice.map(requestFields),
    c_providerReasoningEffort: topLevel,
    continue: continued.map(requestFields),
  };
  console.log(JSON.stringify(report, null, 2));
  if (process.argv.includes("--save-baseline")) {
    const file = new URL("./fixtures/answer-effort-switch-20261013.json", import.meta.url);
    writeFileSync(file, `${JSON.stringify({
      fixtureProvenance: "node node_modules/tsx/dist/cli.mjs tests/answer-effort-probe.mts --dry-run --save-baseline（开关尚未读取环境变量时的整份请求体）",
      natal: natal,
      continue: continued,
    }, null, 2)}\n`);
  }
}

const invoked = process.argv[1]
  ? new URL(`file:///${process.argv[1].replace(/\\/g, "/")}`).href === import.meta.url
  : false;
if (invoked) {
  main().catch((error) => {
    console.error(error);
    process.exitCode = 1;
  });
}