BPOS-MVP/spike/mastra-reopen/src/route-a-dowhile.ts

/**
 * 路线 A:.dowhile() 子图
 *
 * 把「A4 → A5 (→ A6)」整体包成 dowhile 的循环体,循环条件 =「本轮发生了回跳」。
 * 拓扑(对齐种子 unified-quote):
 *   start → A4(定价,活动) → A5(审批)
 *     E11  A5 --verdict=pass--> A6(议价)
 *     E10  A5 --verdict=reject--> A4        (reopen)
 *     E13  A6 --outcome=won--> end
 *     E12  A6 --outcome=revise--> A4        (reopen)
 *
 * 每个活动/审批点都是:起任务 → suspend();终态由外部 run.resume() 送达。
 */
import { Mastra } from '@mastra/core';
import { RequestContext } from '@mastra/core/request-context';
import { InMemoryStore } from '@mastra/core/storage';
import { createStep, createWorkflow, createWorkflowStateReader } from '@mastra/core/workflows';
import { z } from 'zod';
import { banner, DecisionScript, Ledger, type Decision } from './ledger.js';
import { resetTrace, t } from './trace.js';

/** 循环体的输入/输出必须同构 —— dowhile 把上一轮的输出直接当下一轮的输入 */
const RoundState = z.object({
  round: z.number(),
  reopen: z.boolean(),
  reopenedBy: z.string().nullable(),
  price: z.number().nullable(),
  verdict: z.string().nullable(),
  finalOutcome: z.string().nullable(),
  a4TaskId: z.string().nullable(),
});
type RoundState = z.infer<typeof RoundState>;

/** 每个场景一份新的 ledger / script,用模块级变量喂给 step(step 定义是静态的) */
let ledger = new Ledger();
let script = new DecisionScript([]);

/* ------------------------------------------------------------------ 活动点 */

const a4 = createStep({
  id: 'A4',
  description: '定价(活动点)',
  inputSchema: RoundState,
  outputSchema: RoundState,
  resumeSchema: z.object({ price: z.number() }),
  suspendSchema: z.object({ nodeId: z.string(), taskId: z.string(), instanceNo: z.number() }),
  execute: async ({ inputData, resumeData, suspend, runId }) => {
    t('A4.execute', `enter round=${inputData.round} resume=${!!resumeData}`);
    if (!resumeData) {
      const row = ledger.openInstance({
        nodeId: 'A4',
        runId,
        iteration: inputData.round,
        reopenedBy: inputData.reopenedBy,
      });
      t('A4.execute', `起任务 → suspend`, { instanceNo: row.instanceNo, taskId: row.taskId });
      return await suspend(
        { nodeId: 'A4', taskId: row.taskId, instanceNo: row.instanceNo },
        { resumeLabel: `A4-round-${inputData.round}` },
      );
    }
    const taskId = ledger.rows.filter((r) => r.nodeId === 'A4').at(-1)!.taskId;
    ledger.closeTask(taskId, `price=${resumeData.price}`);
    t('A4.execute', `任务终态到达 price=${resumeData.price}`);
    return { ...inputData, price: resumeData.price, a4TaskId: taskId };
  },
});

const a5 = createStep({
  id: 'A5',
  description: '审批(审批点)',
  inputSchema: RoundState,
  outputSchema: RoundState,
  resumeSchema: z.object({ verdict: z.enum(['pass', 'reject']) }),
  suspendSchema: z.object({ nodeId: z.string(), taskId: z.string(), instanceNo: z.number() }),
  execute: async ({ inputData, resumeData, suspend, runId }) => {
    t('A5.execute', `enter round=${inputData.round} resume=${!!resumeData}`);
    if (!resumeData) {
      const row = ledger.openInstance({ nodeId: 'A5', runId, iteration: inputData.round });
      t('A5.execute', `起任务 → suspend`, { instanceNo: row.instanceNo, taskId: row.taskId });
      return await suspend(
        { nodeId: 'A5', taskId: row.taskId, instanceNo: row.instanceNo },
        { resumeLabel: `A5-round-${inputData.round}` },
      );
    }
    const taskId = ledger.rows.filter((r) => r.nodeId === 'A5').at(-1)!.taskId;
    ledger.closeTask(taskId, resumeData.verdict);
    t('A5.execute', `任务终态到达 verdict=${resumeData.verdict}`);
    return { ...inputData, verdict: resumeData.verdict };
  },
});

const a6 = createStep({
  id: 'A6',
  description: '议价(活动点)',
  inputSchema: RoundState,
  outputSchema: RoundState,
  resumeSchema: z.object({ outcome: z.enum(['won', 'revise']) }),
  suspendSchema: z.object({ nodeId: z.string(), taskId: z.string(), instanceNo: z.number() }),
  execute: async ({ inputData, resumeData, suspend, runId }) => {
    t('A6.execute', `enter round=${inputData.round} resume=${!!resumeData}`);
    if (!resumeData) {
      const row = ledger.openInstance({ nodeId: 'A6', runId, iteration: inputData.round });
      t('A6.execute', `起任务 → suspend`, { instanceNo: row.instanceNo, taskId: row.taskId });
      return await suspend(
        { nodeId: 'A6', taskId: row.taskId, instanceNo: row.instanceNo },
        { resumeLabel: `A6-round-${inputData.round}` },
      );
    }
    const taskId = ledger.rows.filter((r) => r.nodeId === 'A6').at(-1)!.taskId;
    ledger.closeTask(taskId, resumeData.outcome);
    t('A6.execute', `任务终态到达 outcome=${resumeData.outcome}`);
    return { ...inputData, finalOutcome: resumeData.outcome };
  },
});

/* --------------------------------------------------------------- 边(出边) */

const e10Reopen = createStep({
  id: 'E10-reopen-A4',
  description: 'A5 --verdict=reject--> A4 (reopen)',
  inputSchema: RoundState,
  outputSchema: RoundState,
  execute: async ({ inputData }) => {
    t('E10.execute', `回边命中,轮次 ${inputData.round} → ${inputData.round + 1}`);
    return { ...inputData, round: inputData.round + 1, reopen: true, reopenedBy: 'E10' };
  },
});

const e12Reopen = createStep({
  id: 'E12-reopen-A4',
  description: 'A6 --outcome=revise--> A4 (reopen)',
  inputSchema: RoundState,
  outputSchema: RoundState,
  execute: async ({ inputData }) => {
    t('E12.execute', `回边命中,轮次 ${inputData.round} → ${inputData.round + 1}`);
    return { ...inputData, round: inputData.round + 1, reopen: true, reopenedBy: 'E12' };
  },
});

const e13End = createStep({
  id: 'E13-end',
  description: 'A6 --outcome=won--> end',
  inputSchema: RoundState,
  outputSchema: RoundState,
  execute: async ({ inputData }) => {
    t('E13.execute', '成交,走向 end');
    return { ...inputData, reopen: false, reopenedBy: null };
  },
});

/**
 * branch 的输出是 { [每一条分支的 stepId]: output | undefined } —— 没命中的分支键也在,值是
 * undefined。所以收口步骤的 inputSchema 必须允许 undefined,并自己做 XOR 判定。
 */
const pick = (id: string) =>
  createStep({
    id,
    inputSchema: z.record(z.string(), RoundState.optional()),
    outputSchema: RoundState,
    execute: async ({ inputData }) => {
      const allKeys = Object.keys(inputData);
      const hit = allKeys.filter((k) => inputData[k] !== undefined);
      t(`${id}.execute`, `branch 收口`, { '所有分支键': allKeys, '命中': hit });
      if (hit.length !== 1) throw new Error(`XOR 违约:命中 ${hit.length} 条出边`);
      return inputData[hit[0]!]!;
    },
  });

/* ------------------------------------------------------------ A6 子图 + 循环体 */

const a6Sub = createWorkflow({
  id: 'A6-sub',
  inputSchema: RoundState,
  outputSchema: RoundState,
})
  .then(a6)
  .branch([
    [
      async ({ inputData, getStepResult, requestContext }) => {
        t('cond(E13 won?)', 'evaluated', {
          finalOutcome: inputData.finalOutcome,
          '上一步输出可见': getStepResult(a6)?.finalOutcome ?? null,
          'requestContext.tenant': requestContext?.get?.('tenant') ?? null,
        });
        return inputData.finalOutcome === 'won';
      },
      e13End,
    ],
    [
      async ({ inputData }) => {
        t('cond(E12 revise?)', 'evaluated', { finalOutcome: inputData.finalOutcome });
        return inputData.finalOutcome === 'revise';
      },
      e12Reopen,
    ],
  ])
  .then(pick('pick-A6-out'))
  .commit();

const roundBody = createWorkflow({
  id: 'round-body',
  inputSchema: RoundState,
  outputSchema: RoundState,
})
  .then(a4)
  .then(a5)
  .branch([
    [
      async ({ inputData, getStepResult, requestContext }) => {
        t('cond(E11 pass?)', 'evaluated', {
          verdict: inputData.verdict,
          '上一步输出可见': getStepResult(a5)?.verdict ?? null,
          'requestContext.tenant': requestContext?.get?.('tenant') ?? null,
        });
        return inputData.verdict === 'pass';
      },
      a6Sub,
    ],
    [
      async ({ inputData }) => {
        t('cond(E10 reject?)', 'evaluated', { verdict: inputData.verdict });
        return inputData.verdict === 'reject';
      },
      e10Reopen,
    ],
  ])
  .then(pick('pick-A5-out'))
  .commit();

const start = createStep({
  id: 'start',
  inputSchema: z.object({ quoteId: z.string() }),
  outputSchema: RoundState,
  execute: async ({ inputData }) => {
    t('start.execute', `起单 ${inputData.quoteId}`);
    return {
      round: 1,
      reopen: false,
      reopenedBy: null,
      price: null,
      verdict: null,
      finalOutcome: null,
      a4TaskId: null,
    };
  },
});

const end = createStep({
  id: 'end',
  inputSchema: RoundState,
  outputSchema: z.object({ rounds: z.number(), finalOutcome: z.string().nullable() }),
  execute: async ({ inputData }) => {
    t('end.execute', `收单 rounds=${inputData.round}`);
    return { rounds: inputData.round, finalOutcome: inputData.finalOutcome };
  },
});

export const routeA = createWorkflow({
  id: 'route-a',
  inputSchema: z.object({ quoteId: z.string() }),
  outputSchema: z.object({ rounds: z.number(), finalOutcome: z.string().nullable() }),
})
  .then(start)
  .dowhile(roundBody, async ({ inputData, iterationCount, requestContext, getInitData }) => {
    t('dowhile.condition', 'evaluated', {
      iterationCount,
      reopen: inputData.reopen,
      round: inputData.round,
      'requestContext.tenant': requestContext?.get?.('tenant') ?? null,
      'getInitData()': getInitData(),
      '继续循环?': inputData.reopen === true,
    });
    return inputData.reopen === true;
  })
  .then(end)
  .commit();

export const storeA = new InMemoryStore();
export const mastraA = new Mastra({
  workflows: { routeA },
  storage: storeA,
  logger: false,
});
const mastra = mastraA;

/** 探针需要拿到 ledger/script 的控制权 */
export const _internals = {
  setLedger: (l: Ledger) => (ledger = l),
  setScript: (s: DecisionScript) => (script = s),
  getLedger: () => ledger,
};

/* ----------------------------------------------------------------- 驱动器 */

type NodeIdT = 'A4' | 'A5' | 'A6';

async function drive(scenario: string, decisions: Decision[]): Promise<void> {
  banner(`路线 A / 场景:${scenario}`);
  resetTrace();
  ledger = new Ledger();
  script = new DecisionScript(decisions);

  const wf = mastra.getWorkflow('routeA');
  const run = await wf.createRun();
  console.log(`  runId = ${run.runId}`);

  // 外部注入的上下文:必须是 RequestContext 实例(传普通对象会在持久化时炸,见 FINDINGS)
  const rc = new RequestContext();
  rc.set('tenant', 'acme');

  let res = await run.start({
    inputData: { quoteId: 'Q-2026-0001' },
    requestContext: rc as never,
  });

  let guard = 0;
  while (res.status === 'suspended') {
    if (++guard > 20) throw new Error('driver: too many suspends');
    const state = await wf.getWorkflowRunById(run.runId);
    const reader = createWorkflowStateReader(state!);
    const s = reader.getSuspendedStep();
    t('[driver]', 'getSuspendedStep()', {
      stepId: s?.stepId,
      path: s?.path,
      executionPath: s?.executionPath,
      suspendPayload: s?.suspendPayload,
      resumeLabels: Object.keys(reader.getResumeLabels()),
    });
    t('[driver]', 'result.suspended (来自 start/resume 的返回值)', res.suspended);

    const nodeId = (s!.suspendPayload as { nodeId: NodeIdT }).nodeId;
    const d = script.next(nodeId);
    const resumeData =
      d.node === 'A4'
        ? { price: d.price }
        : d.node === 'A5'
          ? { verdict: d.verdict }
          : { outcome: d.outcome };
    t('[driver]', `外部推进 ${nodeId}`, resumeData);
    res = await run.resume({ step: s!.path, resumeData });
  }

  console.log(`\n  最终 status = ${res.status}`);
  if (res.status === 'success') console.log('  result =', res.result);
  else console.log('  非 success:', JSON.stringify(res).slice(0, 800));
  ledger.dump(scenario);

  // 跑完之后还能不能查到历史轮次?
  const finalState = await wf.getWorkflowRunById(run.runId);
  console.log('\n  ── 终态快照 steps 的 key ──');
  console.log('  ' + JSON.stringify(Object.keys(finalState?.steps ?? {})));
  const a4Snap = finalState?.steps?.['round-body.A4'] ?? finalState?.steps?.['A4'];
  console.log('  快照里 A4 这一格是数组吗?', Array.isArray(a4Snap), '→', JSON.stringify(a4Snap));
}

// 被 probe 引入时不自动跑三个场景
if (!process.env.SPIKE_NO_AUTORUN) {
  await drive('① 一次过(A4→A5 pass→A6 won→end)', [
    { node: 'A4', price: 100 },
    { node: 'A5', verdict: 'pass' },
    { node: 'A6', outcome: 'won' },
  ]);

  await drive('② A5 驳回一次后再通过', [
    { node: 'A4', price: 100 },
    { node: 'A5', verdict: 'reject' },
    { node: 'A4', price: 90 },
    { node: 'A5', verdict: 'pass' },
    { node: 'A6', outcome: 'won' },
  ]);

  await drive('③ A6 revise 一次后再成交', [
    { node: 'A4', price: 100 },
    { node: 'A5', verdict: 'pass' },
    { node: 'A6', outcome: 'revise' },
    { node: 'A4', price: 80 },
    { node: 'A5', verdict: 'pass' },
    { node: 'A6', outcome: 'won' },
  ]);
}