lam-agent-vicinae/src/__tests__/thread.test.ts

import { describe, expect, it } from "vitest";
import { buildThread, describeActivity, type LiveState } from "../lib/thread";
import type { ChatTurn } from "../lib/types";

const idle: LiveState = { isStreaming: false, response: "", thinking: "" };

const activity = { label: "Reading a page", count: 1 };

function turn(id: string, prompt: string, response: string, overrides: Partial<ChatTurn> = {}): ChatTurn {
  return {
    id,
    prompt,
    response,
    thinking: "",
    model: "m",
    createdAt: 100,
    toolCalls: [],
    attachments: [],
    ...overrides,
  };
}

describe("buildThread", () => {
  it("returns nothing for an empty conversation", () => {
    expect(buildThread([], idle)).toEqual([]);
  });

  it("puts the answer above its question", () => {
    const thread = buildThread([turn("t1", "hello", "hi there")], idle);

    expect(thread.map((message) => message.role)).toEqual(["assistant", "user"]);
    expect(thread[0].text).toBe("hi there");
    expect(thread[1].text).toBe("hello");
  });

  it("shows the newest message first, regardless of who sent it", () => {
    // The search bar is at the top of the window, so the latest answer should be
    // the first thing under it rather than below the fold.
    const thread = buildThread(
      [turn("t1", "first question", "first answer"), turn("t2", "second question", "second answer")],
      idle,
    );

    expect(thread.map((message) => message.text)).toEqual([
      "second answer",
      "second question",
      "first answer",
      "first question",
    ]);
  });

  it("alternates assistant then user", () => {
    const thread = buildThread(
      [turn("t1", "first question", "first answer"), turn("t2", "second question", "second answer")],
      idle,
    );

    expect(thread.map((message) => message.role)).toEqual([
      "assistant",
      "user",
      "assistant",
      "user",
    ]);
  });

  it("gives every message a unique id", () => {
    const thread = buildThread([turn("t1", "a", "b"), turn("t2", "c", "d")], idle);
    const ids = thread.map((message) => message.id);

    expect(new Set(ids).size).toBe(ids.length);
  });

  it("streams the newest turn only, and puts it first", () => {
    const live: LiveState = { isStreaming: true, response: "streaming now", thinking: "hmm" };
    const thread = buildThread([turn("t1", "old", "old answer"), turn("t2", "new", "")], live);

    // The in-flight answer leads...
    expect(thread[0].role).toBe("assistant");
    expect(thread[0].isStreaming).toBe(true);
    expect(thread[0].text).toBe("streaming now");
    expect(thread[0].thinking).toBe("hmm");

    expect(thread[1].role).toBe("user");
    expect(thread[1].text).toBe("new");
    expect(thread[1].isStreaming).toBe(false);

    // ...and the older turn follows, untouched.
    expect(thread[2].text).toBe("old answer");
    expect(thread[3].text).toBe("old");
    expect(thread[2].isStreaming).toBe(false);
  });

  it("carries the activity line while a streaming answer has no text yet", () => {
    const live: LiveState = { isStreaming: true, response: "", thinking: "", status: "Searching the web" };
    const thread = buildThread([turn("t1", "q", "")], live);

    expect(thread[0].status).toBe("Searching the web");
  });

  it("drops the activity line once text has arrived", () => {
    // The placeholder it feeds is only rendered when there is nothing to read.
    const live: LiveState = {
      isStreaming: true,
      response: "partial",
      thinking: "",
      status: "Searching the web",
    };
    const thread = buildThread([turn("t1", "q", "")], live);

    expect(thread[0].status).toBeUndefined();
  });

  it("does not carry an activity line on a finished turn", () => {
    const live: LiveState = { isStreaming: true, response: "", thinking: "", status: "Searching the web" };
    const thread = buildThread([turn("t1", "old", "done"), turn("t2", "q", "")], live);

    // Index 2 is the older, finished exchange's answer.
    expect(thread[2].status).toBeUndefined();
  });

  it("carries tool calls onto the assistant message", () => {
    const withTools = turn("t1", "q", "a", {
      toolCalls: [{ name: "web_search", arguments: {}, result: "ok", isError: false }],
    });
    const thread = buildThread([withTools], idle);

    expect(thread[0].toolCalls).toHaveLength(1);
    expect(thread[0].toolCalls[0].name).toBe("web_search");
    expect(thread[1].toolCalls).toEqual([]);
  });

  it("carries a stored error onto a finished turn", () => {
    const thread = buildThread([turn("t1", "q", "", { error: "Request failed" })], idle);

    expect(thread[0].error).toBe("Request failed");
  });
});

describe("carrying the running tool", () => {
  const live: LiveState = { isStreaming: true, response: "", thinking: "", activity };

  it("puts it on the message being streamed", () => {
    const thread = buildThread([turn("t1", "q", "a")], live);
    const assistant = thread.find((message) => message.role === "assistant");

    expect(assistant?.activity).toEqual(activity);
  });

  it("keeps it off every earlier message", () => {
    // Only the newest turn can be mid-flight, and an older answer carrying a
    // "Searching the web" line would be a claim about a finished turn.
    const thread = buildThread([turn("t1", "q", "a"), turn("t2", "q2", "a2")], live);

    expect(thread.slice(2).every((message) => message.activity === undefined)).toBe(true);
  });

  it("drops it once the turn is over", () => {
    const thread = buildThread([turn("t1", "q", "a")], { ...idle, activity });

    expect(thread.every((message) => message.activity === undefined)).toBe(true);
  });

  it("never puts one on a question", () => {
    const thread = buildThread([turn("t1", "q", "a")], live);

    expect(thread.find((message) => message.role === "user")?.activity).toBeUndefined();
  });
});

describe("describeActivity", () => {
  it("is just the label for a first call", () => {
    expect(describeActivity({ label: "Reading a page", count: 1 })).toBe("Reading a page");
  });

  it("says a repeat is a repeat, because a bare number reads as progress", () => {
    expect(describeActivity({ label: "Reading a page", count: 3 })).toBe(
      "Reading a page · the same call again (3)",
    );
  });
});