pi-context-manager/src/ledger.ts

export type MemoryKind = "goal" | "constraint" | "decision" | "activeTask" | "openQuestion" | "relevantFile";
export type MemoryScope = "branch" | "session";
export type MemorySourceType = "user" | "assistant" | "toolResult" | "compaction" | "branchSummary";

export interface MemoryCandidate {
  kind: MemoryKind;
  subject: string;
  text: string;
  scope: MemoryScope;
  sourceEntryId: string;
  sourceType: MemorySourceType;
  timestamp: number;
  confidence: number;
}

export interface MemoryItem extends MemoryCandidate {
  id: string;
  freshness: number;
  active: boolean;
  supersedesId?: string;
}

export interface LedgerState {
  items: MemoryItem[];
  rollingSummary: string;
}

type MemorySlot = Pick<MemoryCandidate, "kind" | "scope" | "subject">;

export function createEmptyLedger(): LedgerState {
  return { items: [], rollingSummary: "" };
}

function createId(candidate: MemoryCandidate): string {
  return `${candidate.kind}:${candidate.scope}:${candidate.subject}:${candidate.timestamp}`;
}

function ensureUniqueId(items: Pick<MemoryItem, "id">[], baseId: string): string {
  let id = baseId;
  let suffix = 2;

  while (items.some((item) => item.id === id)) {
    id = `${baseId}:${suffix}`;
    suffix += 1;
  }

  return id;
}

function sameSlot(left: MemorySlot, right: MemorySlot) {
  return left.kind === right.kind && left.scope === right.scope && left.subject === right.subject;
}

function createSlotKey(slot: MemorySlot): string {
  return `${slot.kind}\u0000${slot.scope}\u0000${slot.subject}`;
}

function compareStrings(left: string, right: string): number {
  if (left === right) {
    return 0;
  }

  return left < right ? -1 : 1;
}

function compareSameTimestampCandidates(
  left: Pick<MemoryCandidate, "text" | "sourceType" | "sourceEntryId" | "confidence">,
  right: Pick<MemoryCandidate, "text" | "sourceType" | "sourceEntryId" | "confidence">
): number {
  // Exact-timestamp ties should resolve the same way no matter which candidate is processed first.
  const textComparison = compareStrings(left.text, right.text);
  if (textComparison !== 0) {
    return textComparison;
  }

  const sourceTypeComparison = compareStrings(left.sourceType, right.sourceType);
  if (sourceTypeComparison !== 0) {
    return sourceTypeComparison;
  }

  const sourceEntryIdComparison = compareStrings(left.sourceEntryId, right.sourceEntryId);
  if (sourceEntryIdComparison !== 0) {
    return sourceEntryIdComparison;
  }

  if (left.confidence !== right.confidence) {
    return left.confidence > right.confidence ? -1 : 1;
  }

  return 0;
}

function candidateSupersedesPrevious(candidate: MemoryCandidate, previous?: MemoryItem): boolean {
  if (!previous) {
    return true;
  }

  if (candidate.timestamp !== previous.timestamp) {
    return candidate.timestamp > previous.timestamp;
  }

  return compareSameTimestampCandidates(candidate, previous) < 0;
}

function compareSlotItems(left: MemoryItem, right: MemoryItem): number {
  if (left.timestamp !== right.timestamp) {
    return right.timestamp - left.timestamp;
  }

  return compareSameTimestampCandidates(left, right);
}

function normalizeSlotItems(items: MemoryItem[], slot: MemorySlot): MemoryItem[] {
  const slotIndices: number[] = [];
  const slotItems: MemoryItem[] = [];

  items.forEach((item, index) => {
    if (!sameSlot(item, slot)) {
      return;
    }

    slotIndices.push(index);
    slotItems.push(item);
  });

  if (slotItems.length <= 1) {
    return items;
  }

  const sortedSlotItems = [...slotItems].sort(compareSlotItems);
  const slotIds = new Map<string, number>();
  const sortedSlotItemsWithIds = sortedSlotItems.map((item) => {
    const baseId = createId(item);
    const nextSlotIdCount = (slotIds.get(baseId) ?? 0) + 1;
    slotIds.set(baseId, nextSlotIdCount);

    return {
      item,
      id: nextSlotIdCount === 1 ? baseId : `${baseId}:${nextSlotIdCount}`,
    };
  });

  const normalizedSlotItems = sortedSlotItemsWithIds.map(({ item, id }, index) => ({
    ...item,
    id,
    freshness: index === 0 ? item.timestamp : sortedSlotItemsWithIds[index - 1]!.item.timestamp,
    active: index === 0,
    supersedesId: sortedSlotItemsWithIds[index + 1]?.id,
  }));

  const normalizedItems = [...items];
  slotIndices.forEach((slotIndex, index) => {
    normalizedItems[slotIndex] = normalizedSlotItems[index]!;
  });

  return normalizedItems;
}

export function mergeCandidates(state: LedgerState, candidates: MemoryCandidate[]): LedgerState {
  let items = [...state.items];
  const affectedSlots = new Map<string, MemorySlot>();

  for (const candidate of candidates) {
    const previousIndex = items.findIndex((item) => item.active && sameSlot(item, candidate));
    const previous = previousIndex === -1 ? undefined : items[previousIndex];
    const candidateIsNewest = candidateSupersedesPrevious(candidate, previous);

    if (previous && candidateIsNewest) {
      items[previousIndex] = { ...previous, active: false, freshness: candidate.timestamp };
    }

    items.push({
      ...candidate,
      id: ensureUniqueId(items, createId(candidate)),
      freshness: candidate.timestamp,
      active: candidateIsNewest,
      supersedesId: candidateIsNewest ? previous?.id : undefined,
    });

    affectedSlots.set(createSlotKey(candidate), {
      kind: candidate.kind,
      scope: candidate.scope,
      subject: candidate.subject,
    });
  }

  for (const slot of affectedSlots.values()) {
    items = normalizeSlotItems(items, slot);
  }

  return { ...state, items };
}

export function getActiveItems(state: LedgerState, kind?: MemoryKind): MemoryItem[] {
  return state.items.filter((item) => item.active && (kind ? item.kind === kind : true));
}