vpp-ai-platform/packages/services/src/lineage.ts

import { AsyncLocalStorage } from 'node:async_hooks'
import type { BidPayload, BidProposalDraft, DispatchPlanPayload, DispatchProposalDraft, LineageRef, Proposal, ToolCallRef } from '@vpp/domain'
import { BidPayload as BidPayloadSchema, Curve96, DecimalString, DispatchPlanPayload as DispatchPlanPayloadSchema } from '@vpp/domain'
import { proposalDigest } from './digest.js'
import type { SnapshotStore } from './snapshot.js'

/**
 * Lineage recorder (docs/09 §4): every skill call records
 * {tool, version, inputs_ref, outputs_ref}; the buffer travels with the
 * workflow run via AsyncLocalStorage and is frozen into the Proposal.
 */
export class LineageRecorder {
  private readonly als = new AsyncLocalStorage<{ calls: ToolCallRef[]; dataRefs: Set<string> }>()

  run<T>(fn: () => T): T {
    return this.als.run({ calls: [], dataRefs: new Set() }, fn)
  }

  record(call: ToolCallRef): void {
    const store = this.als.getStore()
    if (!store) throw new Error('lineage.record called outside a lineage.run scope')
    store.calls.push(call)
  }

  addDataRef(ref: string): void {
    this.als.getStore()?.dataRefs.add(ref)
  }

  current(): { tool_calls: ToolCallRef[]; data_refs: string[] } {
    const store = this.als.getStore()
    if (!store) throw new Error('lineage.current called outside a lineage.run scope')
    return { tool_calls: [...store.calls], data_refs: [...store.dataRefs].sort() }
  }
}

/** Wrap any deterministic skill call so it is snapshotted and recorded. */
export function withLineage<I, O>(
  lineage: LineageRecorder,
  snapshots: SnapshotStore,
  def: { tool: string; version: string; idGen: () => string },
  call: (input: I) => O,
): (input: I) => { output: O; call: ToolCallRef } {
  return (input: I) => {
    const inputs_ref = snapshots.put(input)
    const output = call(input)
    const outputs_ref = snapshots.put(output)
    const ref: ToolCallRef = { tool_call_id: def.idGen(), tool: def.tool, version: def.version, inputs_ref, outputs_ref }
    lineage.record(ref)
    return { output, call: ref }
  }
}

export class LineageRefError extends Error {
  constructor(message: string) {
    super(message)
    this.name = 'LineageRefError'
  }
}

/** Dereference {tool_call_id, path} into the recorded output snapshot. */
export function resolveRef(ref: LineageRef, calls: ToolCallRef[], snapshots: SnapshotStore): unknown {
  const call = calls.find((c) => c.tool_call_id === ref.tool_call_id)
  if (!call) throw new LineageRefError(`tool call ${ref.tool_call_id} is not in this proposal's lineage`)
  const output = snapshots.get(call.outputs_ref)
  if (output === undefined) throw new LineageRefError(`output snapshot for ${ref.tool_call_id} missing`)
  let node: unknown = output
  for (const part of ref.path.split('.')) {
    if (node === null || typeof node !== 'object') throw new LineageRefError(`path ${ref.path} does not resolve in ${ref.tool_call_id}`)
    node = Array.isArray(node) ? node[Number(part)] : (node as Record<string, unknown>)[part]
    if (node === undefined) throw new LineageRefError(`path ${ref.path} does not resolve in ${ref.tool_call_id}`)
  }
  return node
}

/**
 * The P2 assembler (docs/09 §3): builds a Proposal from an LLM draft that
 * contains only references. Numbers are dereferenced from lineage snapshots;
 * a draft with a literal number is rejected by the schema before we get here.
 */
export function assembleBidProposal(
  draft: BidProposalDraft,
  ctx: {
    id: string
    originator: Proposal['originator']
    lineage: { tool_calls: ToolCallRef[]; data_refs: string[] }
    ledger_version: number
    policy_pack_version: string
    snapshots: SnapshotStore
    now: string
  },
): Proposal {
  const prices = Curve96.parse(resolveRef(draft.prices_ref, ctx.lineage.tool_calls, ctx.snapshots))
  const quantities = Curve96.parse(resolveRef(draft.quantities_ref, ctx.lineage.tool_calls, ctx.snapshots))
  const revenue = DecimalString.parse(resolveRef(draft.expected_revenue_ref, ctx.lineage.tool_calls, ctx.snapshots))
  const payload: BidPayload = BidPayloadSchema.parse({
    kind: 'BID',
    market_date: draft.market_date,
    prices_yuan_per_mwh: prices,
    quantities_mwh: quantities,
    expected_revenue_yuan: revenue,
  })
  const lineage = {
    tool_calls: ctx.lineage.tool_calls,
    data_refs: ctx.lineage.data_refs,
    ledger_version: ctx.ledger_version,
    policy_pack_version: ctx.policy_pack_version,
  }
  const base = {
    id: ctx.id,
    type: 'BID' as const,
    timescale: 'DAY_AHEAD' as const,
    originator: ctx.originator,
    payload,
    lineage,
    envelope_ref: null,
    status: 'DRAFT' as const,
    created_at: ctx.now,
  }
  return { ...base, digest: proposalDigest(base) }
}

/** DISPATCH_PLAN assembler: same discipline, resource-agent originator, EVENT trigger. */
export function assembleDispatchProposal(
  draft: DispatchProposalDraft,
  ctx: {
    id: string
    award_ref: string
    originator: Proposal['originator']
    lineage: { tool_calls: ToolCallRef[]; data_refs: string[] }
    ledger_version: number
    policy_pack_version: string
    snapshots: SnapshotStore
    now: string
  },
): Proposal {
  const payload: DispatchPlanPayload = DispatchPlanPayloadSchema.parse({
    kind: 'DISPATCH_PLAN',
    market_date: draft.market_date,
    award_ref: ctx.award_ref,
    total_target_mw: Curve96.parse(resolveRef(draft.total_target_ref, ctx.lineage.tool_calls, ctx.snapshots)),
    allocations: resolveRef(draft.allocations_ref, ctx.lineage.tool_calls, ctx.snapshots),
  })
  const base = {
    id: ctx.id,
    type: 'DISPATCH_PLAN' as const,
    timescale: 'DAY_AHEAD' as const,
    originator: ctx.originator,
    payload,
    lineage: { tool_calls: ctx.lineage.tool_calls, data_refs: ctx.lineage.data_refs, ledger_version: ctx.ledger_version, policy_pack_version: ctx.policy_pack_version },
    envelope_ref: null,
    status: 'DRAFT' as const,
    created_at: ctx.now,
  }
  return { ...base, digest: proposalDigest(base) }
}