《互联网产品的 AI 架构·多能力编排与任务闭环》4.2 节点间的上下文预算分配

inputCap = maxContextTokens − outputReserveTokens
floor    = fixedSystemTokens + sharedSliceTokens + minPrivateTokens

// file: src/orch/context-budget.ts
// tier: L2
// verify: node --test test/orch/context-budget.test.ts

import { AiError } from "../core/errors.ts";

/**
 * 一个节点对上下文的四段需求：
 * 固定系统提示、共享片段、节点私有工作集，以及必须留给输出的预留。
 * 前三段计入输入预算，第四段从模型窗口里扣。
 */
export interface NodeContextDemand {
  id: string;
  /** 数值越大越先保底、越先增长。 */
  priority: number;
  fixedSystemTokens: number;
  sharedSliceTokens: number;
  desiredPrivateTokens: number;
  minPrivateTokens: number;
  outputReserveTokens: number;
  maxContextTokens: number;
}

export interface ContextBudgetSpec {
  /** 所有节点输入 token 的总额；不包含输出预留。 */
  totalInputTokens: number;
  nodes: readonly NodeContextDemand[];
}

export type DropReason = "budget-floor" | "context-window";

export interface NodeAllocation {
  id: string;
  priority: number;
  allocated: boolean;
  systemTokens: number;
  sharedTokens: number;
  privateTokens: number;
  outputReserveTokens: number;
  inputTokens: number;
  /** 输入 + 输出预留，必须不超过节点的模型窗口。 */
  contextTokens: number;
  desiredPrivateTokens: number;
  /** 相对 desiredPrivateTokens 少拿到的私有 token。 */
  clippedPrivateTokens: number;
  dropReason?: DropReason;
}

export interface BudgetAllocation {
  totalInputTokens: number;
  usedInputTokens: number;
  remainingInputTokens: number;
  allocations: readonly NodeAllocation[];
  dropped: readonly { id: string; reason: DropReason }[];
  /** 被保留节点少拿到的私有 token 总数；被丢弃节点不计入。 */
  clippedPrivateTokens: number;
  feasibleCount: number;
  fullySatisfied: boolean;
}

function fail(message: string): never {
  throw new AiError("invalid-request", `context-budget: ${message}`);
}

function integer(value: number, field: string, min = 0): number {
  if (!Number.isInteger(value) || value < min) fail(`${field} must be an integer >= ${min}`);
  return value;
}

function text(value: string, field: string): string {
  if (typeof value !== "string" || value.trim() === "") fail(`${field} must be a non-empty string`);
  return value;
}

function validate(spec: ContextBudgetSpec): void {
  integer(spec.totalInputTokens, "totalInputTokens");
  if (spec.nodes.length === 0) fail("至少需要一个节点");
  const seen = new Set<string>();
  for (const node of spec.nodes) {
    text(node.id, "node.id");
    if (seen.has(node.id)) fail(`重复的节点 id ${node.id}`);
    seen.add(node.id);
    integer(node.priority, `${node.id}.priority`);
    integer(node.fixedSystemTokens, `${node.id}.fixedSystemTokens`);
    integer(node.sharedSliceTokens, `${node.id}.sharedSliceTokens`);
    integer(node.desiredPrivateTokens, `${node.id}.desiredPrivateTokens`);
    integer(node.minPrivateTokens, `${node.id}.minPrivateTokens`);
    integer(node.outputReserveTokens, `${node.id}.outputReserveTokens`);
    integer(node.maxContextTokens, `${node.id}.maxContextTokens`, 1);
    if (node.minPrivateTokens > node.desiredPrivateTokens) {
      fail(`${node.id}.minPrivateTokens 不能大于 desiredPrivateTokens`);
    }
    if (node.outputReserveTokens > node.maxContextTokens) {
      fail(`${node.id}.outputReserveTokens 不能超过 maxContextTokens`);
    }
  }
}

function drop(node: NodeContextDemand, reason: DropReason): NodeAllocation {
  return {
    id: node.id,
    priority: node.priority,
    allocated: false,
    systemTokens: 0,
    sharedTokens: 0,
    privateTokens: 0,
    outputReserveTokens: 0,
    inputTokens: 0,
    contextTokens: 0,
    desiredPrivateTokens: node.desiredPrivateTokens,
    clippedPrivateTokens: 0,
    dropReason: reason,
  };
}

/**
 * 两阶段分配：先按优先级保底，再按同一顺序把剩余预算长到 desired。
 * 保底失败或超出模型窗口的节点被丢弃，不会得到半截上下文；
 * 保底成功后私有部分可以被裁到 min 与 desired 之间。
 */
export function allocateContextBudget(spec: ContextBudgetSpec): BudgetAllocation {
  validate(spec);
  const ordered = [...spec.nodes].sort((a, b) => {
    if (a.priority !== b.priority) return b.priority - a.priority;
    return a.id < b.id ? -1 : a.id > b.id ? 1 : 0;
  });

  let remaining = spec.totalInputTokens;
  const dropped: { id: string; reason: DropReason }[] = [];
  const kept: { node: NodeContextDemand; inputCap: number; floor: number }[] = [];

  for (const node of ordered) {
    const inputCap = node.maxContextTokens - node.outputReserveTokens;
    const floor = node.fixedSystemTokens + node.sharedSliceTokens + node.minPrivateTokens;
    if (floor > inputCap) {
      dropped.push({ id: node.id, reason: "context-window" });
      continue;
    }
    if (floor > remaining) {
      dropped.push({ id: node.id, reason: "budget-floor" });
      continue;
    }
    remaining -= floor;
    kept.push({ node, inputCap, floor });
  }

  const allocations: NodeAllocation[] = [];
  let clipped = 0;
  let satisfied = true;

  for (const item of kept) {
    const { node, inputCap } = item;
    const maxPrivate = Math.min(
      node.desiredPrivateTokens,
      inputCap - node.fixedSystemTokens - node.sharedSliceTokens,
    );
    const wanted = maxPrivate - node.minPrivateTokens;
    const take = Math.max(0, Math.min(wanted, remaining));
    remaining -= take;
    const privateTokens = node.minPrivateTokens + take;
    const inputTokens = node.fixedSystemTokens + node.sharedSliceTokens + privateTokens;
    const clippedPrivate = node.desiredPrivateTokens - privateTokens;
    if (clippedPrivate > 0) satisfied = false;
    clipped += clippedPrivate;
    allocations.push({
      id: node.id,
      priority: node.priority,
      allocated: true,
      systemTokens: node.fixedSystemTokens,
      sharedTokens: node.sharedSliceTokens,
      privateTokens,
      outputReserveTokens: node.outputReserveTokens,
      inputTokens,
      contextTokens: inputTokens + node.outputReserveTokens,
      desiredPrivateTokens: node.desiredPrivateTokens,
      clippedPrivateTokens: clippedPrivate,
    });
  }

  for (const item of dropped) {
    const node = ordered.find((candidate) => candidate.id === item.id);
    if (node) allocations.push(drop(node, item.reason));
  }

  return {
    totalInputTokens: spec.totalInputTokens,
    usedInputTokens: spec.totalInputTokens - remaining,
    remainingInputTokens: remaining,
    allocations,
    dropped,
    clippedPrivateTokens: clipped,
    feasibleCount: kept.length,
    fullySatisfied: satisfied,
  };
}