《互联网产品的 AI 架构·推理性能与成本工程》4.2 流式输出的提前终止

// file: src/perf/output-stream.ts
// tier: L1
// verify: node --test test/perf/output-stream.test.ts

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

export interface EarlyStopPolicy {
  /** 少于这么多输出 token 时，即使置信度很高也不停止。 */
  minTokens: number;
  /** 输出预算硬上限，由 4.1 的 planOutputBudget 给出。 */
  maxTokens: number;
  deadlineMs: number;
  confidenceThreshold: number;
  /** 首帧之后失败不重试：停止后的结果就是最终结果。 */
  noRetryAfterFirstFrame: boolean;
  /** 结构化输出必须等结构闭合后才允许按置信度停止。 */
  requireClosure?: boolean;
}

export interface EarlyStopState {
  tokens: number;
  confidence: number;
  firstFrameAtMs: number | null;
  elapsedMs: number;
  closed: boolean;
  cancelled: boolean;
  cancelReason: string | null;
  stopped: boolean;
}

export type EarlyStopReason = "continue" | "confidence" | "max-tokens" | "deadline" | "user-cancel";

export interface EarlyStopDecision {
  stop: boolean;
  reason: EarlyStopReason;
  /** 已经生成、必须计量的 token。 */
  committedTokens: number;
  /** 停止时最多还能省下的 token，等于预算上限减已生成。 */
  savedTokensUpperBound: number;
}

/** 与卷一取消路径一致：CJK 约 1 token，ASCII 约 4 字符 1 token。 */
export function approximateOutputTokens(text: string): number {
  let tokens = 0;
  for (const char of text) tokens += (char.codePointAt(0) ?? 0) > 0x7f ? 1 : 0.25;
  return Math.ceil(tokens);
}

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

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

function requireConfidence(value: number, field: string): number {
  if (typeof value !== "number" || !Number.isFinite(value) || value < 0 || value > 1) {
    fail(`${field} must be a finite number in [0, 1]`);
  }
  return value;
}

export interface EarlyStopOptions {
  policy: EarlyStopPolicy;
  /** 注入计数器，测试里可以直接按 token 计数。 */
  countTokens?: (text: string) => number;
}

/** 只负责“是否停、为什么停、欠了多少 token”，关闭连接交给卷一的取消信号。 */
export class EarlyStopController {
  readonly #policy: EarlyStopPolicy;
  readonly #count: (text: string) => number;
  #tokens = 0;
  #confidence = 0;
  #firstFrameAtMs: number | null = null;
  #elapsedMs = 0;
  #closed = false;
  #cancelled = false;
  #cancelReason: string | null = null;
  #decision: EarlyStopDecision | null = null;

  constructor(options: EarlyStopOptions) {
    const policy = options.policy;
    requireInteger(policy.minTokens, "minTokens", 0);
    requireInteger(policy.maxTokens, "maxTokens", 1);
    requireInteger(policy.deadlineMs, "deadlineMs", 1);
    requireConfidence(policy.confidenceThreshold, "confidenceThreshold");
    if (policy.minTokens > policy.maxTokens) fail("minTokens must not exceed maxTokens");
    this.#policy = { ...policy };
    this.#count = options.countTokens ?? approximateOutputTokens;
  }

  get state(): EarlyStopState {
    return {
      tokens: this.#tokens, confidence: this.#confidence,
      firstFrameAtMs: this.#firstFrameAtMs, elapsedMs: this.#elapsedMs,
      closed: this.#closed, cancelled: this.#cancelled,
      cancelReason: this.#cancelReason, stopped: this.#decision?.stop === true,
    };
  }

  markClosed(): void {
    this.#closed = true;
  }

  cancel(reason = "user"): EarlyStopDecision {
    this.#cancelled = true;
    this.#cancelReason = reason;
    return this.#finalize("user-cancel");
  }

  onDelta(text: string, confidence: number, elapsedMs: number): EarlyStopDecision {
    if (this.#decision !== null && this.#decision.stop) return this.#decision;
    requireConfidence(confidence, "confidence");
    this.#elapsedMs = requireInteger(elapsedMs, "elapsedMs", 0);
    if (text !== "") {
      this.#tokens += this.#count(text);
      this.#confidence = confidence;
      this.#firstFrameAtMs ??= elapsedMs;
    }
    if (this.#cancelled) return this.#finalize("user-cancel");
    if (this.#elapsedMs >= this.#policy.deadlineMs) return this.#finalize("deadline");
    if (this.#tokens >= this.#policy.maxTokens) return this.#finalize("max-tokens");

    const confidenceReached = this.#tokens >= this.#policy.minTokens
      && this.#confidence >= this.#policy.confidenceThreshold;
    const closureReady = this.#policy.requireClosure !== true || this.#closed;
    if (confidenceReached && closureReady) return this.#finalize("confidence");

    return {
      stop: false, reason: "continue", committedTokens: this.#tokens,
      savedTokensUpperBound: this.#policy.maxTokens - this.#tokens,
    };
  }

  #finalize(reason: Exclude<EarlyStopReason, "continue">): EarlyStopDecision {
    if (this.#decision !== null && this.#decision.stop) return this.#decision;
    this.#decision = {
      stop: true, reason, committedTokens: this.#tokens,
      savedTokensUpperBound: Math.max(0, this.#policy.maxTokens - this.#tokens),
    };
    return this.#decision;
  }
}

export interface StopEconomicsInput {
  emittedTokens: number;
  expectedFullTokens: number;
  confidence: number;
  microsPerToken: number;
  /** 错误停止后重新问一次的成本，包含输入与重建输出。 */
  reaskCostMicros: number;
}

export interface StopEconomics {
  savedTokens: number;
  savedMicros: number;
  wrongProbability: number;
  expectedLossMicros: number;
  netMicros: number;
  breakEvenConfidence: number;
  worthStopping: boolean;
}

/** 省下的是“本来还会生成的 token”，赔的是“停早了要重问”。 */
export function evaluateStopEconomics(input: StopEconomicsInput): StopEconomics {
  const emitted = requireInteger(input.emittedTokens, "emittedTokens", 0);
  const full = requireInteger(input.expectedFullTokens, "expectedFullTokens", 0);
  const micros = requireInteger(input.microsPerToken, "microsPerToken", 0);
  const reask = requireInteger(input.reaskCostMicros, "reaskCostMicros", 0);
  const confidence = requireConfidence(input.confidence, "confidence");
  const savedTokens = Math.max(0, full - emitted);
  const savedMicros = savedTokens * micros;
  const wrongProbability = 1 - confidence;
  const expectedLossMicros = Math.ceil(wrongProbability * reask - 1e-9);
  const rawBreakEven = reask === 0 ? 0 : Math.min(1, Math.max(0, 1 - savedMicros / reask));

  return {
    savedTokens, savedMicros, wrongProbability, expectedLossMicros,
    netMicros: savedMicros - expectedLossMicros,
    breakEvenConfidence: Math.round(rawBreakEven * 1_000_000) / 1_000_000,
    worthStopping: savedMicros - expectedLossMicros > 0,
  };
}

export interface StreamFailureInput {
  firstFrameSent: boolean;
  errorRetryable: boolean;
  noRetryAfterFirstFrame: boolean;
  committedTokens: number;
}

export interface StreamFailureDecision {
  retryable: boolean;
  committedTokens: number;
  reason: "retry-allowed" | "no-retry-after-first-frame" | "error-not-retryable";
}

/** 首帧之后失败不重试；已生成的 token 仍然计费。 */
export function decideStreamFailure(input: StreamFailureInput): StreamFailureDecision {
  const committedTokens = requireInteger(input.committedTokens, "committedTokens", 0);
  if (input.firstFrameSent && input.noRetryAfterFirstFrame) {
    return { retryable: false, committedTokens, reason: "no-retry-after-first-frame" };
  }
  if (!input.errorRetryable) {
    return { retryable: false, committedTokens, reason: "error-not-retryable" };
  }
  return { retryable: true, committedTokens, reason: "retry-allowed" };
}

node --test test/perf/output-stream.test.ts

ℹ tests 12
ℹ pass 12
ℹ fail 0
ℹ duration_ms 158.712417