《互联网产品的 AI 架构·多能力编排与任务闭环》4.3 长任务的会话状态与记忆边界

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

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

/** 会话存储的一条记录；tenantId + sessionId 才是唯一键。 */
export interface SessionEntry {
  tenantId: string;
  sessionId: string;
  tokens: number;
  priority: number;
  pinned: boolean;
  createdAt: number;
  lastAccessAt: number;
  /** 绝对过期时刻；与 idleTtlMs 取更早者。 */
  expiresAt: number;
}

export interface SessionInput {
  tenantId: string;
  sessionId: string;
  tokens: number;
  priority: number;
  pinned?: boolean;
  createdAt?: number;
  expiresAt?: number;
}

export type SessionEvictionPolicy = "lru" | "priority";

export interface SessionPolicy {
  idleTtlMs: number;
  absoluteTtlMs: number;
  maxEntries: number;
  maxTokens: number;
  maxTokensPerSession: number;
  eviction: SessionEvictionPolicy;
}

export type SessionRemovalReason =
  | "idle-ttl"
  | "absolute-ttl"
  | "capacity-entries"
  | "capacity-tokens";

export interface SessionRemoval {
  tenantId: string;
  sessionId: string;
  reason: SessionRemovalReason;
  tokens: number;
}

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

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

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 validatePolicy(policy: SessionPolicy): void {
  finite(policy.idleTtlMs, "idleTtlMs", 1);
  finite(policy.absoluteTtlMs, "absoluteTtlMs", 1);
  integer(policy.maxEntries, "maxEntries", 1);
  integer(policy.maxTokens, "maxTokens", 1);
  integer(policy.maxTokensPerSession, "maxTokensPerSession", 1);
  if (policy.maxTokensPerSession > policy.maxTokens) {
    fail("maxTokensPerSession 不能超过 maxTokens");
  }
  if (policy.eviction !== "lru" && policy.eviction !== "priority") {
    fail(`unknown eviction policy ${String(policy.eviction)}`);
  }
}

function isBetterVictim(
  candidateKey: string,
  candidate: SessionEntry,
  currentKey: string,
  current: SessionEntry,
  policy: SessionPolicy,
): boolean {
  if (candidate.pinned !== current.pinned) return !candidate.pinned;
  if (policy.eviction === "priority" && candidate.priority !== current.priority) {
    return candidate.priority < current.priority;
  }
  if (candidate.lastAccessAt !== current.lastAccessAt) {
    return candidate.lastAccessAt < current.lastAccessAt;
  }
  return candidateKey < currentKey;
}

/**
 * 最小会话存储：滑动空闲 TTL、绝对 TTL、条数与 token 容量、跨租户隔离。
 * 每次 get 都会刷新 lastAccessAt；put 覆盖时不会延长已有的绝对过期时刻。
 */
export class SessionStore {
  private readonly policy: SessionPolicy;
  private readonly sessions = new Map<string, SessionEntry>();
  private readonly removalLog: SessionRemoval[] = [];
  private tokenTotal = 0;

  constructor(policy: SessionPolicy) {
    validatePolicy(policy);
    this.policy = { ...policy };
  }

  get size(): number {
    return this.sessions.size;
  }

  get totalTokens(): number {
    return this.tokenTotal;
  }

  get removals(): readonly SessionRemoval[] {
    return this.removalLog.slice();
  }

  private key(tenantId: string, sessionId: string): string {
    return `${tenantId}\u0000${sessionId}`;
  }

  private drop(key: string, entry: SessionEntry, reason: SessionRemovalReason, log: SessionRemoval[]): void {
    this.sessions.delete(key);
    this.tokenTotal -= entry.tokens;
    log.push({ tenantId: entry.tenantId, sessionId: entry.sessionId, reason, tokens: entry.tokens });
  }

  /** 不刷新访问时间，用于观测与测试。 */
  peek(tenantId: string, sessionId: string): SessionEntry | undefined {
    const entry = this.sessions.get(this.key(tenantId, sessionId));
    return entry === undefined ? undefined : { ...entry };
  }

  get(tenantId: string, sessionId: string, now: number): SessionEntry | undefined {
    text(tenantId, "tenantId");
    text(sessionId, "sessionId");
    finite(now, "now");
    this.sweep(now);
    const entry = this.sessions.get(this.key(tenantId, sessionId));
    if (entry === undefined || entry.tenantId !== tenantId) return undefined;
    entry.lastAccessAt = now;
    return { ...entry };
  }

  put(input: SessionInput, now: number): SessionEntry {
    text(input.tenantId, "tenantId");
    text(input.sessionId, "sessionId");
    integer(input.tokens, "tokens");
    integer(input.priority, "priority");
    finite(now, "now");
    if (input.tokens > this.policy.maxTokensPerSession) {
      fail(`${input.sessionId}.tokens 超过单会话上限 ${this.policy.maxTokensPerSession}`);
    }
    this.sweep(now);
    const key = this.key(input.tenantId, input.sessionId);
    const existing = this.sessions.get(key);
    const createdAt = existing?.createdAt ?? input.createdAt ?? now;
    const requestedExpiry = input.expiresAt ?? now + this.policy.absoluteTtlMs;
    const expiresAt = existing === undefined
      ? requestedExpiry
      : Math.min(existing.expiresAt, requestedExpiry);
    if (expiresAt <= now) fail(`${input.sessionId} 在写入时已过期`);

    const entry: SessionEntry = {
      tenantId: input.tenantId,
      sessionId: input.sessionId,
      tokens: input.tokens,
      priority: input.priority,
      pinned: input.pinned === true,
      createdAt,
      lastAccessAt: now,
      expiresAt,
    };
    this.tokenTotal += input.tokens - (existing?.tokens ?? 0);
    this.sessions.set(key, entry);
    this.enforceCapacity();
    return { ...entry };
  }

  /**
   * 追加记忆。超过单会话 token 上限时直接拒绝，而不是悄悄截断或无限增长。
   */
  append(tenantId: string, sessionId: string, deltaTokens: number, now: number): boolean {
    text(tenantId, "tenantId");
    text(sessionId, "sessionId");
    integer(deltaTokens, "deltaTokens", 1);
    finite(now, "now");
    this.sweep(now);
    const key = this.key(tenantId, sessionId);
    const entry = this.sessions.get(key);
    if (entry === undefined || entry.tenantId !== tenantId) return false;
    const nextTokens = entry.tokens + deltaTokens;
    if (nextTokens > this.policy.maxTokensPerSession) {
      fail(`${sessionId} 追加后超过单会话上限 ${this.policy.maxTokensPerSession}`);
    }
    entry.tokens = nextTokens;
    entry.lastAccessAt = now;
    this.tokenTotal += deltaTokens;
    this.enforceCapacity();
    return true;
  }

  /** 清理空闲过期与绝对过期的记录。 */
  sweep(now: number): readonly SessionRemoval[] {
    finite(now, "now");
    const removed: SessionRemoval[] = [];
    for (const [key, entry] of this.sessions) {
      const absoluteExpired = now >= entry.expiresAt;
      const idleExpired = now - entry.lastAccessAt >= this.policy.idleTtlMs;
      if (absoluteExpired) {
        this.drop(key, entry, "absolute-ttl", removed);
      } else if (idleExpired) {
        this.drop(key, entry, "idle-ttl", removed);
      }
    }
    this.removalLog.push(...removed);
    return removed;
  }

  private selectVictim(): { key: string; entry: SessionEntry } | undefined {
    let best: { key: string; entry: SessionEntry } | undefined;
    for (const [key, entry] of this.sessions) {
      if (best === undefined || isBetterVictim(key, entry, best.key, best.entry, this.policy)) {
        best = { key, entry };
      }
    }
    return best;
  }

  private enforceCapacity(): void {
    while (this.sessions.size > this.policy.maxEntries || this.tokenTotal > this.policy.maxTokens) {
      const victim = this.selectVictim();
      if (victim === undefined) break;
      const reason: SessionRemovalReason = this.sessions.size > this.policy.maxEntries
        ? "capacity-entries"
        : "capacity-tokens";
      this.drop(victim.key, victim.entry, reason, this.removalLog);
    }
  }
}