BPOS-TMP-REPO/spike/mastra-reopen/src/route-b-external.ts

/**
 * 路线 B:外部胶合(回跳 = 结束本 run,外部用同一份 workflow 定义新起一个 run)
 *
 * workflow 定义里 **没有循环**,就是一张 DAG:
 *   A4 → A5 --E11 pass--> A6 --E13 won--> end
 *              --E10 reject--> [回边出口]
 *                        A6 --E12 revise--> [回边出口]
 * 「回边出口」是一个终止步骤:它把 reopen 意图写进 workflow 的 result,run 正常 success 结束。
 * 外部胶合层看到 result.reopen 就 createRun 一个新 run,把重开载荷(重开边、实例号、上一个 runId)
 * 当 inputData 传进去。run 之间用 runId 链 + resourceId 关联。
 */
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';

/** 一个 run = 一趟「A4 起,到 end 或到回边出口」 */
const RunIn = z.object({
  quoteId: z.string(),
  /** A4 的第几个实例(= 第几次重开 + 1),由外部胶合层递增后传进来 */
  instanceSeq: z.number(),
  /** 由哪条回边重开;首开为 null */
  reopenedBy: z.string().nullable(),
  /** 上一个 run 的 runId,形成链 */
  prevRunId: z.string().nullable(),
});
type RunIn = z.infer<typeof RunIn>;

const Carry = RunIn.extend({
  price: z.number().nullable(),
  verdict: z.string().nullable(),
  finalOutcome: z.string().nullable(),
});

const RunOut = z.object({
  /** 本 run 是否以回跳结束 */
  reopen: z.boolean(),
  /** 回跳边号 */
  reopenedBy: z.string().nullable(),
  /** 回跳目标节点 */
  reopenTarget: z.string().nullable(),
  finalOutcome: z.string().nullable(),
  instanceSeq: z.number(),
});

let ledger = new Ledger();
let script = new DecisionScript([]);

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

const a4 = createStep({
  id: 'A4',
  inputSchema: RunIn,
  outputSchema: Carry,
  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 seq=${inputData.instanceSeq} resume=${!!resumeData}`);
    if (!resumeData) {
      const row = ledger.openInstance({
        nodeId: 'A4',
        runId,
        iteration: inputData.instanceSeq,
        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' },
      );
    }
    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, verdict: null, finalOutcome: null };
  },
});

const a5 = createStep({
  id: 'A5',
  inputSchema: Carry,
  outputSchema: Carry,
  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 seq=${inputData.instanceSeq} resume=${!!resumeData}`);
    if (!resumeData) {
      const row = ledger.openInstance({ nodeId: 'A5', runId, iteration: inputData.instanceSeq });
      t('A5.execute', '起任务 → suspend', { instanceNo: row.instanceNo, taskId: row.taskId });
      return await suspend(
        { nodeId: 'A5', taskId: row.taskId, instanceNo: row.instanceNo },
        { resumeLabel: 'A5' },
      );
    }
    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',
  inputSchema: Carry,
  outputSchema: Carry,
  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 seq=${inputData.instanceSeq} resume=${!!resumeData}`);
    if (!resumeData) {
      const row = ledger.openInstance({ nodeId: 'A6', runId, iteration: inputData.instanceSeq });
      t('A6.execute', '起任务 → suspend', { instanceNo: row.instanceNo, taskId: row.taskId });
      return await suspend(
        { nodeId: 'A6', taskId: row.taskId, instanceNo: row.instanceNo },
        { resumeLabel: 'A6' },
      );
    }
    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 reopenExit = (id: string, edge: string) =>
  createStep({
    id,
    inputSchema: Carry,
    outputSchema: RunOut,
    execute: async ({ inputData }) => {
      t(`${id}.execute`, `回边命中 → 本 run 到此为止,交给外部胶合层重开 A4`);
      return {
        reopen: true,
        reopenedBy: edge,
        reopenTarget: 'A4',
        finalOutcome: inputData.finalOutcome,
        instanceSeq: inputData.instanceSeq,
      };
    },
  });

const e10Exit = reopenExit('E10-reopen-exit', 'E10');
const e12Exit = reopenExit('E12-reopen-exit', 'E12');

const e13End = createStep({
  id: 'E13-end',
  inputSchema: Carry,
  outputSchema: RunOut,
  execute: async ({ inputData }) => {
    t('E13.execute', '成交,走向 end');
    return {
      reopen: false,
      reopenedBy: null,
      reopenTarget: null,
      finalOutcome: inputData.finalOutcome,
      instanceSeq: inputData.instanceSeq,
    };
  },
});

const pick = (id: string) =>
  createStep({
    id,
    inputSchema: z.record(z.string(), RunOut.optional()),
    outputSchema: RunOut,
    execute: async ({ inputData }) => {
      const hit = Object.keys(inputData).filter((k) => inputData[k] !== undefined);
      t(`${id}.execute`, 'branch 收口', { '所有分支键': Object.keys(inputData), '命中': hit });
      if (hit.length !== 1) throw new Error(`XOR 违约:命中 ${hit.length} 条出边`);
      return inputData[hit[0]!]!;
    },
  });

const a6Sub = createWorkflow({ id: 'A6-sub', inputSchema: Carry, outputSchema: RunOut })
  .then(a6)
  .branch([
    [
      async ({ inputData, requestContext }) => {
        t('cond(E13 won?)', 'evaluated', {
          finalOutcome: inputData.finalOutcome,
          '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';
      },
      e12Exit,
    ],
  ])
  .then(pick('pick-A6-out'))
  .commit();

export const routeB = createWorkflow({
  id: 'route-b',
  inputSchema: RunIn,
  outputSchema: RunOut,
})
  .then(a4)
  .then(a5)
  .branch([
    [
      async ({ inputData, requestContext }) => {
        t('cond(E11 pass?)', 'evaluated', {
          verdict: inputData.verdict,
          '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';
      },
      e10Exit,
    ],
  ])
  .then(pick('pick-A5-out'))
  .commit();

export const storeB = new InMemoryStore();
export const mastraB = new Mastra({
  workflows: { routeB },
  storage: storeB,
  logger: false,
});
const mastra = mastraB;

export const _internals = {
  setLedger: (l: Ledger) => (ledger = l),
  setScript: (s: DecisionScript) => (script = s),
  getLedger: () => ledger,
};

/* ------------------------------------------------------- 外部胶合层(驱动器) */

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

/** run 链:外部胶合层自己维护的关联表(主工程里就是「实例表 + 边留痕」) */
type ChainRow = { seq: number; runId: string; prevRunId: string | null; reopenedBy: string | null; endedAs: string };

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

  const wf = mastra.getWorkflow('routeB');
  const quoteId = `Q-${Math.random().toString(36).slice(2, 8)}`;
  const chain: ChainRow[] = [];

  let input: RunIn = { quoteId, instanceSeq: 1, reopenedBy: null, prevRunId: null };
  let guard = 0;

  for (;;) {
    if (++guard > 10) throw new Error('driver: too many reopens');

    // 自己指定 runId —— 可读、可幂等、可与实例号一一对应
    const runId = `${quoteId}-inst-${input.instanceSeq}`;
    const run = await wf.createRun({ runId, resourceId: quoteId });
    t('[glue]', `createRun`, { runId: run.runId, resourceId: quoteId, inputData: input });

    const rc = new RequestContext();
    rc.set('tenant', 'acme');

    let res = await run.start({ inputData: input, requestContext: rc as never });

    let g2 = 0;
    while (res.status === 'suspended') {
      if (++g2 > 10) throw new Error('driver: too many suspends');
      const state = await wf.getWorkflowRunById(run.runId);
      const reader = createWorkflowStateReader(state!);
      const s = reader.getSuspendedStep();
      t('[glue]', 'getSuspendedStep()', {
        stepId: s?.stepId,
        path: s?.path,
        suspendPayload: s?.suspendPayload,
        resumeLabels: Object.keys(reader.getResumeLabels()),
      });
      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('[glue]', `外部推进 ${nodeId}`, resumeData);
      res = await run.resume({ step: s!.path, resumeData });
    }

    if (res.status !== 'success') {
      console.log('  非 success:', JSON.stringify(res).slice(0, 600));
      return;
    }
    const out = res.result;
    chain.push({
      seq: input.instanceSeq,
      runId: run.runId,
      prevRunId: input.prevRunId,
      reopenedBy: input.reopenedBy,
      endedAs: out.reopen ? `reopen(${out.reopenedBy}) → ${out.reopenTarget}` : `end(${out.finalOutcome})`,
    });

    if (!out.reopen) {
      t('[glue]', `run 结束:${JSON.stringify(out)} —— 流程完结`);
      break;
    }
    t('[glue]', `run 结束:回跳 ${out.reopenedBy} → ${out.reopenTarget},新起一个 run`);
    input = {
      quoteId,
      instanceSeq: input.instanceSeq + 1,
      reopenedBy: out.reopenedBy,
      prevRunId: run.runId,
    };
  }

  console.log('\n  ── run 链 ──');
  for (const c of chain) {
    console.log(
      `  seq=${c.seq}\trunId=${c.runId}\tprev=${c.prevRunId ?? '-'}\treopenedBy=${c.reopenedBy ?? '-'}\t${c.endedAs}`,
    );
  }
  ledger.dump(scenario);

  // 旧 run 留痕:流程都跑完了,还能不能查到第 1 个 run 的快照?
  console.log('\n  ── 历史 run 是否还查得到(旧实例留痕不动) ──');
  for (const c of chain) {
    const st = await wf.getWorkflowRunById(c.runId);
    const a4 = st?.steps?.['A4'] as { output?: { price?: number }; suspendPayload?: { taskId?: string } } | undefined;
    console.log(
      `  ${c.runId}: status=${st?.status} A4.taskId=${a4?.suspendPayload?.taskId} A4.price=${a4?.output?.price} result=${JSON.stringify(st?.result)}`,
    );
  }
  const runs = await wf.listWorkflowRuns({ resourceId: quoteId });
  console.log(`  listWorkflowRuns({resourceId:'${quoteId}'}) → total=${runs.total}, runIds=${JSON.stringify(runs.runs.map((r) => r.runId))}`);
}

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