clanker-agent/packages/coding-agent/test/session-manager/labels.test.ts

import { describe, expect, it } from "vitest";
import {
  type LabelEntry,
  SessionManager,
} from "../../src/core/session-manager.js";

describe("SessionManager labels", () => {
  it("sets and gets labels", () => {
    const session = SessionManager.inMemory();

    const msgId = session.appendMessage({
      role: "user",
      content: "hello",
      timestamp: 1,
    });

    // No label initially
    expect(session.getLabel(msgId)).toBeUndefined();

    // Set a label
    const labelId = session.appendLabelChange(msgId, "checkpoint");
    expect(session.getLabel(msgId)).toBe("checkpoint");

    // Label entry should be in entries
    const entries = session.getEntries();
    const labelEntry = entries.find((e) => e.type === "label") as LabelEntry;
    expect(labelEntry).toBeDefined();
    expect(labelEntry.id).toBe(labelId);
    expect(labelEntry.targetId).toBe(msgId);
    expect(labelEntry.label).toBe("checkpoint");
  });

  it("clears labels with undefined", () => {
    const session = SessionManager.inMemory();

    const msgId = session.appendMessage({
      role: "user",
      content: "hello",
      timestamp: 1,
    });

    session.appendLabelChange(msgId, "checkpoint");
    expect(session.getLabel(msgId)).toBe("checkpoint");

    // Clear the label
    session.appendLabelChange(msgId, undefined);
    expect(session.getLabel(msgId)).toBeUndefined();
  });

  it("last label wins", () => {
    const session = SessionManager.inMemory();

    const msgId = session.appendMessage({
      role: "user",
      content: "hello",
      timestamp: 1,
    });

    session.appendLabelChange(msgId, "first");
    session.appendLabelChange(msgId, "second");
    session.appendLabelChange(msgId, "third");

    expect(session.getLabel(msgId)).toBe("third");
  });

  it("labels are included in tree nodes", () => {
    const session = SessionManager.inMemory();

    const msg1Id = session.appendMessage({
      role: "user",
      content: "hello",
      timestamp: 1,
    });
    const msg2Id = session.appendMessage({
      role: "assistant",
      content: [{ type: "text", text: "hi" }],
      api: "anthropic-messages",
      provider: "anthropic",
      model: "test",
      usage: {
        input: 1,
        output: 1,
        cacheRead: 0,
        cacheWrite: 0,
        totalTokens: 2,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
      },
      stopReason: "stop",
      timestamp: 2,
    });

    session.appendLabelChange(msg1Id, "start");
    session.appendLabelChange(msg2Id, "response");

    const tree = session.getTree();

    // Find the message nodes (skip label entries)
    const msg1Node = tree.find((n) => n.entry.id === msg1Id);
    expect(msg1Node?.label).toBe("start");

    // msg2 is a child of msg1
    const msg2Node = msg1Node?.children.find((n) => n.entry.id === msg2Id);
    expect(msg2Node?.label).toBe("response");
  });

  it("labels are preserved in createBranchedSession", () => {
    const session = SessionManager.inMemory();

    const msg1Id = session.appendMessage({
      role: "user",
      content: "hello",
      timestamp: 1,
    });
    const msg2Id = session.appendMessage({
      role: "assistant",
      content: [{ type: "text", text: "hi" }],
      api: "anthropic-messages",
      provider: "anthropic",
      model: "test",
      usage: {
        input: 1,
        output: 1,
        cacheRead: 0,
        cacheWrite: 0,
        totalTokens: 2,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
      },
      stopReason: "stop",
      timestamp: 2,
    });

    session.appendLabelChange(msg1Id, "important");
    session.appendLabelChange(msg2Id, "also-important");

    // Branch from msg2 (in-memory mode returns null, but updates internal state)
    session.createBranchedSession(msg2Id);

    // Labels should be preserved
    expect(session.getLabel(msg1Id)).toBe("important");
    expect(session.getLabel(msg2Id)).toBe("also-important");

    // New label entries should exist
    const entries = session.getEntries();
    const labelEntries = entries.filter(
      (e) => e.type === "label",
    ) as LabelEntry[];
    expect(labelEntries).toHaveLength(2);
  });

  it("labels not on path are not preserved in createBranchedSession", () => {
    const session = SessionManager.inMemory();

    const msg1Id = session.appendMessage({
      role: "user",
      content: "hello",
      timestamp: 1,
    });
    const msg2Id = session.appendMessage({
      role: "assistant",
      content: [{ type: "text", text: "hi" }],
      api: "anthropic-messages",
      provider: "anthropic",
      model: "test",
      usage: {
        input: 1,
        output: 1,
        cacheRead: 0,
        cacheWrite: 0,
        totalTokens: 2,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
      },
      stopReason: "stop",
      timestamp: 2,
    });
    const msg3Id = session.appendMessage({
      role: "user",
      content: "followup",
      timestamp: 3,
    });

    // Label all messages
    session.appendLabelChange(msg1Id, "first");
    session.appendLabelChange(msg2Id, "second");
    session.appendLabelChange(msg3Id, "third");

    // Branch from msg2 (excludes msg3)
    session.createBranchedSession(msg2Id);

    // Only labels for msg1 and msg2 should be preserved
    expect(session.getLabel(msg1Id)).toBe("first");
    expect(session.getLabel(msg2Id)).toBe("second");
    expect(session.getLabel(msg3Id)).toBeUndefined();
  });

  it("labels are not included in buildSessionContext", () => {
    const session = SessionManager.inMemory();

    const msgId = session.appendMessage({
      role: "user",
      content: "hello",
      timestamp: 1,
    });
    session.appendLabelChange(msgId, "checkpoint");

    const ctx = session.buildSessionContext();
    expect(ctx.messages).toHaveLength(1);
    expect(ctx.messages[0].role).toBe("user");
  });

  it("throws when labeling non-existent entry", () => {
    const session = SessionManager.inMemory();

    expect(() => session.appendLabelChange("non-existent", "label")).toThrow(
      "Entry non-existent not found",
    );
  });
});