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

import test from "node:test";
import assert from "node:assert/strict";
import { SegmentIndex } from "../src/segment-index.ts";

const assistant = {
  role: "assistant",
  content: [{ type: "toolCall", id: "call-1", name: "bash", arguments: { command: "npm test" } }],
  timestamp: 10,
};
const result = {
  role: "toolResult",
  toolCallId: "call-1",
  toolName: "bash",
  content: [{ type: "text", text: `failure\n${"x".repeat(12_000)}\nstack tail` }],
  isError: true,
  timestamp: 11,
};

// Regression: a ~14.5K-char read result projected under a 1M window must reach
// the model byte-for-byte unchanged — masking only happens via context_edit
// drafts under real pressure, never on the hot projection path.
test("projection is byte-for-byte identity for a 300-line read result under a 1M window", () => {
  const readBody = Array.from({ length: 300 }, (_unused, index) => `line ${index + 1}: ${"y".repeat(38)}`).join("\n");
  const readResult = {
    role: "toolResult",
    toolCallId: "call-read",
    toolName: "read",
    content: [{ type: "text", text: readBody }],
    timestamp: 12,
  };
  assert.ok(readBody.length > 14_000 && readBody.length < 15_000, `fixture ${readBody.length} chars`);

  const index = new SegmentIndex("session-1");
  index.hydrate();
  const snapshot = index.project([assistant, readResult], {
    environmentHash: "env",
    environmentTokens: 100,
    contextWindow: 1_000_000,
    maskedResults: 0,
    reclaimedTokens: 0,
  });

  assert.strictEqual(snapshot.messages[1], readResult);
  assert.strictEqual(snapshot.messages[1]!.content[0].text, readBody);
  assert.ok(snapshot.messages[1]!.content[0].text.includes("line 150:"));
  assert.equal(snapshot.reclaimedTokens, 0);
  assert.equal(snapshot.maskedResults, 0);
});

test("projection caches identical input and handles append-only deltas", () => {
  const index = new SegmentIndex("session-2");
  index.hydrate();
  const messages = [assistant, result];
  const params = { environmentHash: "env", environmentTokens: 100, contextWindow: 300_000, maskedResults: 0, reclaimedTokens: 0 };
  const first = index.project(messages, params);
  const second = index.project(messages, params);
  assert.strictEqual(first, second);
  assert.ok(first.rawTokens > 0);

  const appended = index.project([...messages, { role: "user", content: "next", timestamp: 13 }], params);
  assert.notStrictEqual(appended, first);
  assert.equal(appended.messages.length, 3);
  assert.strictEqual(appended.messages[1], result);
});

test("changed masking stats publish a new snapshot without mutating the old one", () => {
  const index = new SegmentIndex("session-stats");
  index.hydrate();
  const params = {
    environmentHash: "env",
    environmentTokens: 100,
    contextWindow: 300_000,
    maskedResults: 0,
    reclaimedTokens: 0,
  };
  const first = index.project([assistant, result], params);
  const same = index.project([assistant, result], params);
  assert.strictEqual(same, first);

  const updated = index.project([assistant, result], { ...params, maskedResults: 1, reclaimedTokens: 500 });
  assert.notStrictEqual(updated, first);
  assert.equal(updated.maskedResults, 1);
  assert.equal(updated.reclaimedTokens, 500);
  assert.equal(updated.rawTokens, updated.effectiveTokens + 500);
  // The previously published object stays immutable.
  assert.equal(first.maskedResults, 0);
  assert.equal(first.reclaimedTokens, 0);
  // Unchanged stats keep returning the same cached object.
  assert.strictEqual(index.project([assistant, result], { ...params, maskedResults: 1, reclaimedTokens: 500 }), updated);
});

test("usage-backed baseline is adjusted only by projected delta", () => {
  const index = new SegmentIndex("session-3");
  index.hydrate();
  const initial = index.project([assistant], {
    environmentHash: "env",
    environmentTokens: 50,
    contextWindow: 100_000,
    maskedResults: 0,
    reclaimedTokens: 0,
  });
  index.recordUsage(1_000, initial.messageEstimate, "env");
  const next = index.project([assistant, { role: "user", content: "next", timestamp: 12 }], {
    environmentHash: "env",
    environmentTokens: 50,
    contextWindow: 100_000,
    maskedResults: 0,
    reclaimedTokens: 0,
  });
  assert.equal(next.confidence, "usage-backed");
  assert.ok(next.effectiveTokens >= 1_000);
  assert.ok(next.effectiveTokens < 1_100);
});