my-pi/pi-compaction/src/segment-index.ts

import { createHash } from "node:crypto";
import { estimateTokens } from "@earendil-works/pi-coding-agent";
import type { ProjectionSnapshot, UsageBaseline } from "./types.js";

export function timestampOf(value: unknown): number {
  if (typeof value === "number" && Number.isFinite(value)) return value;
  if (typeof value === "string") {
    const parsed = Date.parse(value);
    if (Number.isFinite(parsed)) return parsed;
  }
  return 0;
}

export function contentText(content: unknown): string {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  return content
    .flatMap((part) => part && typeof part === "object" && (part as any).type === "text" && typeof (part as any).text === "string"
      ? [(part as any).text]
      : [])
    .join("\n");
}

function hash(value: string): string {
  return createHash("sha256").update(value).digest("hex").slice(0, 20);
}

function messageKey(message: any): string {
  const role = String(message?.role ?? "unknown");
  const timestamp = timestampOf(message?.timestamp);
  if (role === "toolResult") return `${role}:${message.toolCallId ?? ""}:${timestamp}`;
  if (role === "assistant") {
    const blocks = Array.isArray(message.content) ? message.content : [];
    const ids = blocks.flatMap((block: any) => block?.type === "toolCall" ? [block.id] : []).join(",");
    return `${role}:${timestamp}:${ids}:${blocks.length}`;
  }
  return `${role}:${timestamp}:${typeof message?.content === "string" ? message.content.length : (message?.content?.length ?? 0)}`;
}

/**
 * Incremental context projection: tracks the branch's message list, estimates
 * tokens, and maintains the usage-backed effective-token baseline. Projection
 * is identity — pi-compaction never rewrites messages on this path; context
 * pressure reduction happens through append-only context_edit entries.
 */
export class SegmentIndex {
  private generation = 0;
  private usageBaseline?: UsageBaseline;
  private sourceMessages: any[] = [];
  private sourceKeys: string[] = [];
  private projectedMessages: any[] = [];
  private rawEstimates: number[] = [];
  private effectiveEstimates: number[] = [];
  private cached?: ProjectionSnapshot;

  readonly sessionId: string;

  constructor(sessionId: string) {
    this.sessionId = sessionId;
  }

  hydrate(): void {
    this.reset();
    this.generation++;
  }

  reset(): void {
    this.generation = 0;
    this.usageBaseline = undefined;
    this.sourceMessages = [];
    this.sourceKeys = [];
    this.projectedMessages = [];
    this.rawEstimates = [];
    this.effectiveEstimates = [];
    this.cached = undefined;
  }

  recordUsage(tokens: number, messageEstimate: number, environmentHash: string): void {
    if (!Number.isFinite(tokens) || tokens <= 0) return;
    this.usageBaseline = { tokens, messageEstimate, environmentHash };
    this.cached = undefined;
  }

  project(messages: readonly any[], params: {
    environmentHash: string;
    environmentTokens: number;
    contextWindow: number;
    maskedResults: number;
    reclaimedTokens: number;
  }): ProjectionSnapshot {
    const sameTail = messages.length === this.sourceMessages.length
      && messages.length > 0
      && messageKey(messages[messages.length - 1]) === this.sourceKeys[this.sourceKeys.length - 1];
    if (this.cached && sameTail && this.cached.environmentHash === params.environmentHash) {
      if (this.cached.maskedResults === params.maskedResults
        && this.cached.reclaimedTokens === params.reclaimedTokens) {
        return this.cached;
      }
      // Masking stats changed but the message list did not: publish a new
      // immutable snapshot object instead of re-projecting or mutating.
      this.cached = {
        ...this.cached,
        maskedResults: params.maskedResults,
        reclaimedTokens: params.reclaimedTokens,
        rawTokens: this.cached.effectiveTokens + params.reclaimedTokens,
      };
      return this.cached;
    }

    const appendOnly = this.sourceMessages.length > 0
      && messages.length >= this.sourceMessages.length
      && messageKey(messages[this.sourceMessages.length - 1]) === this.sourceKeys[this.sourceKeys.length - 1]
      && this.cached?.environmentHash === params.environmentHash;

    let start = 0;
    if (!appendOnly) {
      this.sourceMessages = [];
      this.sourceKeys = [];
      this.projectedMessages = [];
      this.rawEstimates = [];
      this.effectiveEstimates = [];
    } else {
      start = this.sourceMessages.length;
    }

    for (let index = start; index < messages.length; index++) {
      const message = messages[index];
      this.sourceMessages.push(message);
      this.sourceKeys.push(messageKey(message));
      this.projectedMessages.push(message);
      this.rawEstimates.push(estimateTokens(message as any));
      this.effectiveEstimates.push(estimateTokens(message as any));
    }

    const effectiveMessageTokens = this.effectiveEstimates.reduce((sum, value) => sum + value, 0);
    const baseline = this.usageBaseline;
    const usageBacked = baseline?.environmentHash === params.environmentHash;
    const effectiveTokens = usageBacked
      ? Math.max(0, baseline.tokens + effectiveMessageTokens - baseline.messageEstimate)
      : effectiveMessageTokens + params.environmentTokens;
    const confidence = usageBacked ? "usage-backed" as const : "estimated" as const;
    const lowerBound = Math.max(0, Math.floor(effectiveTokens * (usageBacked ? 0.98 : 0.9)));
    const upperBound = Math.ceil(effectiveTokens * (usageBacked ? 1.05 : 1.2));
    const fingerprint = hash(`${this.sessionId}:${this.generation}:${params.environmentHash}:${this.sourceKeys.join("|")}`);

    this.cached = {
      sessionId: this.sessionId,
      generation: this.generation,
      fingerprint,
      environmentHash: params.environmentHash,
      messages: this.projectedMessages,
      messageEstimate: effectiveMessageTokens,
      rawTokens: effectiveTokens + params.reclaimedTokens,
      effectiveTokens,
      reclaimedTokens: params.reclaimedTokens,
      contextWindow: params.contextWindow,
      confidence,
      lowerBound,
      upperBound,
      segmentCount: this.sourceMessages.length,
      maskedResults: params.maskedResults,
    };
    return this.cached;
  }

  getSnapshot(): ProjectionSnapshot | undefined {
    return this.cached;
  }
}