pixlstash-lmstudio

import { describe, it } from "node:test";
import assert from "node:assert/strict";
import {
  LEADING_NOISE,
  collectUsedIds,
  makeLeadingNoiseStripper,
} from "../src/predictionLoop.ts";

// Minimal Chat-shaped stub: just what collectUsedIds reads.
const chat = (...texts: string[]) => ({
  getMessagesArray: () => texts.map((t) => ({ getText: () => t })),
});

describe("collectUsedIds", () => {
  it("returns an empty set for an empty history", () => {
    assert.deepEqual([...collectUsedIds(chat())], []);
  });

  it("finds a single id from a markdown image reference", () => {
    const ids = collectUsedIds(chat("![pixlstash 42](pixlstash-42.png)"));
    assert.deepEqual([...ids], [42]);
  });

  it("finds ids across multiple messages and dedupes", () => {
    const ids = collectUsedIds(
      chat(
        "prose then ![](pixlstash-7.webp) more prose",
        "another (pixlstash-42.jpg) reference",
        "and a repeat of pixlstash-7.webp",
      ),
    );
    assert.deepEqual([...ids].sort((a, b) => a - b), [7, 42]);
  });

  it("ignores text that doesn't match the pattern", () => {
    const ids = collectUsedIds(chat("we talked about pixlstash earlier but no files"));
    assert.deepEqual([...ids], []);
  });

  it("requires the `pixlstash-` prefix (not just any number)", () => {
    const ids = collectUsedIds(chat("file-99.png and image-99.png"));
    assert.deepEqual([...ids], []);
  });
});

describe("LEADING_NOISE", () => {
  it("matches the common [TOOL_RESULT] variants", () => {
    assert.match("[TOOL_RESULT]\nhello", LEADING_NOISE);
    assert.match("[tool_result] hi", LEADING_NOISE);
    assert.match("[/TOOL_RESULT]\n", LEADING_NOISE);
    assert.match("  [TOOL RESULT] hi", LEADING_NOISE);
  });

  it("matches Llama-style <|tool_result|> tokens", () => {
    assert.match("<|tool_result|>hi", LEADING_NOISE);
    assert.match("<tool_result>hi", LEADING_NOISE);
    assert.match("</tool_result>hi", LEADING_NOISE);
  });

  it("does not match normal prose", () => {
    assert.doesNotMatch("Hello there!", LEADING_NOISE);
    assert.doesNotMatch("[narration in brackets] not a tool result", LEADING_NOISE);
  });
});

describe("makeLeadingNoiseStripper", () => {
  // Capture every appendText call.
  const capture = () => {
    const chunks: string[] = [];
    return { block: { appendText: (t: string) => void chunks.push(t) }, chunks };
  };

  it("strips a leading [TOOL_RESULT] marker before emitting prose", () => {
    const c = capture();
    const s = makeLeadingNoiseStripper(() => c.block);
    s.emit("[TOOL_RESULT]\nHello, world!");
    s.flush();
    assert.equal(c.chunks.join(""), "Hello, world!");
  });

  it("passes prose through unchanged when no marker is present", () => {
    const c = capture();
    const s = makeLeadingNoiseStripper(() => c.block);
    s.emit("Just narrating ");
    s.emit("normally.");
    s.flush();
    assert.equal(c.chunks.join(""), "Just narrating normally.");
  });

  it("only buffers until the threshold; later fragments stream verbatim", () => {
    const c = capture();
    const s = makeLeadingNoiseStripper(() => c.block);
    s.emit("She walks in.");
    // Past the sentence boundary, the buffer is flushed.
    s.emit(" Then more text.");
    s.flush();
    assert.equal(c.chunks.join(""), "She walks in. Then more text.");
  });

  it("flush() emits whatever is buffered, stripping the marker if present", () => {
    const c = capture();
    const s = makeLeadingNoiseStripper(() => c.block);
    s.emit("[TOOL_RESULT]"); // short, no boundary — kept in buffer
    s.flush(); // forces resolution
    assert.equal(c.chunks.join(""), "");
  });

  it("reset() re-arms stripping for a new text block (e.g. after a tool call)", () => {
    const captures = [capture(), capture()];
    let current = 0;
    const s = makeLeadingNoiseStripper(() => captures[current].block);

    s.emit("[TOOL_RESULT] first block");
    s.flush();
    assert.equal(captures[0].chunks.join(""), "first block");

    // Open a "new" block and reset; the new block's leading marker should also strip.
    current = 1;
    s.reset();
    s.emit("[tool_result]\nsecond block");
    s.flush();
    assert.equal(captures[1].chunks.join(""), "second block");
  });
});