TJWaterAgent/tests/results/store.test.ts

import { afterEach, beforeEach, describe, expect, it } from "bun:test";
import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { ResultReferenceResolver } from "../../src/results/resolver.js";
import {
  RESULT_REFERENCE_KIND,
  RESULT_REFERENCE_SOURCE,
  ResultReferenceStore,
} from "../../src/results/store.js";

describe("ResultReferenceResolver", () => {
  let tempDir: string;
  let store: ResultReferenceStore;
  let resolver: ResultReferenceResolver;

  beforeEach(async () => {
    tempDir = await mkdtemp(join(tmpdir(), "tjwater-result-ref-"));
    store = new ResultReferenceStore(tempDir, 60_000);
    resolver = new ResultReferenceResolver(store);
    await store.initialize();
  });

  afterEach(async () => {
    await rm(tempDir, { force: true, recursive: true });
  });

  it("stores metadata for new referenced results and resolves them", async () => {
    const record = await resolver.register({
      actorKey: "actor-1",
      clientSessionId: "client-1",
      data: [{ id: "J1" }, { id: "J2" }],
      kind: RESULT_REFERENCE_KIND.dynamicHttpResult,
      projectId: "project-1",
      projectKey: "project-key-1",
      schemaVersion: 1,
      sessionId: "session-1",
      source: RESULT_REFERENCE_SOURCE.dynamicHttp,
      traceId: "trace-1",
    });

    expect(record.kind).toBe(RESULT_REFERENCE_KIND.dynamicHttpResult);
    expect(record.schemaVersion).toBe(1);
    expect(record.source).toBe(RESULT_REFERENCE_SOURCE.dynamicHttp);

    const result = await resolver.getAuthorized(
      record.resultRef,
      {
        actorKey: "actor-1",
        projectId: "project-1",
      },
      {
        maxItems: 1,
      },
    );

    expect(result).not.toBeNull();
    expect(result?.kind).toBe(RESULT_REFERENCE_KIND.dynamicHttpResult);
    expect(result?.schema_version).toBe(1);
    expect(result?.source).toBe(RESULT_REFERENCE_SOURCE.dynamicHttp);
    expect(result?.data).toEqual([{ id: "J1" }]);
  });

  it("keeps legacy result refs readable while defaulting metadata", async () => {
    const legacyRef = "res-aaaaaaaaaaaaaaaa";
    await writeFile(
      join(tempDir, `${legacyRef}.json`),
      JSON.stringify(
        {
          resultRef: legacyRef,
          actorKey: "actor-legacy",
          clientSessionId: "client-legacy",
          createdAt: "2026-05-21T00:00:00.000Z",
          data: { nodes: ["J1"] },
          preview: {
            count: 1,
            fields: ["nodes"],
            sample: { nodes: ["J1"] },
            summary: "object<1 fields>",
          },
          projectId: "project-legacy",
          projectKey: "project-key-legacy",
          sessionId: "session-legacy",
          sizeBytes: 16,
          traceId: "trace-legacy",
        },
        null,
        2,
      ),
      "utf8",
    );

    const record = await store.getAuthorizedRecord(legacyRef, {
      actorKey: "actor-legacy",
      projectId: "project-legacy",
    });

    expect(record).not.toBeNull();
    expect(record?.kind).toBe(RESULT_REFERENCE_KIND.dynamicHttpResult);
    expect(record?.schemaVersion).toBe(1);
    expect(record?.source).toBe(RESULT_REFERENCE_SOURCE.legacy);
  });

  it("rejects malformed refs, mismatched kinds, and auth mismatches", async () => {
    const malformedRef = "res-bbbbbbbbbbbbbbbb";
    await writeFile(
      join(tempDir, `${malformedRef}.json`),
      JSON.stringify(
        {
          resultRef: malformedRef,
          createdAt: "2026-05-21T00:00:00.000Z",
          data: { value: 1 },
          preview: {
            count: 1,
            fields: ["value"],
            sample: { value: 1 },
            summary: "object<1 fields>",
          },
          projectId: "project-1",
          projectKey: "project-key-1",
          sessionId: "session-1",
          sizeBytes: 10,
          traceId: "trace-1",
        },
        null,
        2,
      ),
      "utf8",
    );

    const malformed = await store.getAuthorizedRecord(malformedRef, {
      actorKey: "actor-1",
      projectId: "project-1",
    });
    expect(malformed).toBeNull();

    const renderRecord = await resolver.register({
      actorKey: "actor-2",
      clientSessionId: "client-2",
      data: {
        node_area_map: {
          J1: "DMA-1",
        },
      },
      kind: RESULT_REFERENCE_KIND.renderJunctionsPayload,
      projectId: "project-2",
      projectKey: "project-key-2",
      schemaVersion: 1,
      sessionId: "session-2",
      source: RESULT_REFERENCE_SOURCE.agentGenerated,
      traceId: "trace-2",
    });

    const wrongKind = await resolver.getFullAuthorized(
      renderRecord.resultRef,
      {
        actorKey: "actor-2",
        projectId: "project-2",
      },
      {
        expectedKind: RESULT_REFERENCE_KIND.dynamicHttpResult,
      },
    );
    expect(wrongKind).toBeNull();

    const wrongActor = await resolver.getFullAuthorized(renderRecord.resultRef, {
      actorKey: "actor-other",
      projectId: "project-2",
    });
    expect(wrongActor).toBeNull();
  });

  it("registers render refs from local wrapper files and normalizes payloads", async () => {
    const filePath = join(tempDir, "render-wrapper.json");
    await writeFile(
      filePath,
      JSON.stringify(
        {
          metadata: {
            createdAt: "2026-05-21T00:00:00.000Z",
            projectId: "project-3",
          },
          location: {
            file_path: filePath,
          },
          data: {
            node_area_map: {
              J1: "DMA-1",
              J2: 2,
            },
            area_ids: ["DMA-1", " DMA-2 "],
            area_colors: {
              "DMA-1": "#ff0000",
              "DMA-2": "#00ff00",
            },
          },
          createdAt: "2026-05-21T00:00:00.000Z",
        },
        null,
        2,
      ),
      "utf8",
    );

    const record = await resolver.registerRenderPayloadFile(filePath, {
      actorKey: "actor-3",
      clientSessionId: "client-3",
      projectId: "project-3",
      projectKey: "project-key-3",
      sessionId: "session-3",
      source: RESULT_REFERENCE_SOURCE.migration,
      traceId: "trace-3",
    });

    expect(record.kind).toBe(RESULT_REFERENCE_KIND.renderJunctionsPayload);
    expect(record.source).toBe(RESULT_REFERENCE_SOURCE.migration);

    const result = await resolver.getFullAuthorized(
      record.resultRef,
      {
        actorKey: "actor-3",
        projectId: "project-3",
      },
      {
        expectedKind: RESULT_REFERENCE_KIND.renderJunctionsPayload,
      },
    );

    expect(result?.data).toEqual({
      node_area_map: {
        J1: "DMA-1",
        J2: "2",
      },
      area_ids: ["DMA-1", "DMA-2"],
      area_colors: {
        "DMA-1": "#ff0000",
        "DMA-2": "#00ff00",
      },
    });
  });

  it("repairs wrapper files that omit metadata and location", async () => {
    const filePath = join(tempDir, "render-wrapper-missing-fields.json");
    await writeFile(
      filePath,
      JSON.stringify(
        {
          data: {
            node_area_map: {
              J1: "DMA-1",
            },
          },
          createdAt: "2026-05-21T00:00:00.000Z",
        },
        null,
        2,
      ),
      "utf8",
    );

    const record = await resolver.registerRenderPayloadFile(filePath, {
      actorKey: "actor-4",
      clientSessionId: "client-4",
      projectId: "project-4",
      projectKey: "project-key-4",
      sessionId: "session-4",
      source: RESULT_REFERENCE_SOURCE.migration,
      traceId: "trace-4",
    });

    expect(record.kind).toBe(RESULT_REFERENCE_KIND.renderJunctionsPayload);

    const repaired = JSON.parse(await readFile(filePath, "utf8")) as {
      metadata?: Record<string, unknown>;
      location?: Record<string, unknown>;
    };

    expect(repaired.metadata).toEqual({
      createdAt: "2026-05-21T00:00:00.000Z",
      projectId: "project-4",
    });
    expect(repaired.location).toEqual({
      file_path: filePath,
    });
  });

  it("repairs wrapper files whose location points elsewhere", async () => {
    const filePath = join(tempDir, "render-wrapper-wrong-location.json");
    await writeFile(
      filePath,
      JSON.stringify(
        {
          metadata: {
            createdAt: "2026-05-21T00:00:00.000Z",
          },
          location: {
            file_path: "/tmp/elsewhere.json",
            source: "legacy",
          },
          data: {
            node_area_map: {
              J1: "DMA-1",
            },
          },
        },
        null,
        2,
      ),
      "utf8",
    );

    await resolver.registerRenderPayloadFile(filePath, {
        actorKey: "actor-4",
        clientSessionId: "client-4",
        projectId: "project-4",
        projectKey: "project-key-4",
        sessionId: "session-4",
        source: RESULT_REFERENCE_SOURCE.migration,
        traceId: "trace-4",
      });

    const repaired = JSON.parse(await readFile(filePath, "utf8")) as {
      metadata?: Record<string, unknown>;
      location?: Record<string, unknown>;
    };

    expect(repaired.metadata).toEqual({
      createdAt: "2026-05-21T00:00:00.000Z",
      projectId: "project-4",
    });
    expect(repaired.location).toEqual({
      file_path: filePath,
      source: "legacy",
    });
  });

  it("resolves legacy render payload files when callers include the json suffix", async () => {
    const legacyRef = "res-c2fcee33-577e";
    await writeFile(
      join(tempDir, `${legacyRef}.json`),
      JSON.stringify(
        {
          data: {
            node_area_map: {
              J1: "DMA-1",
              J2: 2,
            },
            area_ids: ["DMA-1"],
          },
          createdAt: "2026-05-21T00:00:00.000Z",
          projectId: "project-legacy-render",
        },
        null,
        2,
      ),
      "utf8",
    );

    const result = await resolver.getFullAuthorized(
      `${legacyRef}.json`,
      {
        actorKey: "actor-legacy-render",
        clientSessionId: "chat-legacy-render",
        projectId: "project-legacy-render",
      },
      {
        expectedKind: RESULT_REFERENCE_KIND.renderJunctionsPayload,
      },
    );

    expect(result?.result_ref).toBe(legacyRef);
    expect(result?.kind).toBe(RESULT_REFERENCE_KIND.renderJunctionsPayload);
    expect(result?.source).toBe(RESULT_REFERENCE_SOURCE.legacy);
    expect(result?.data).toEqual({
      node_area_map: {
        J1: "DMA-1",
        J2: "2",
      },
      area_ids: ["DMA-1"],
    });
  });

  it("keeps legacy render payload files scoped to their project", async () => {
    const legacyRef = "res-dddddddddddddddd";
    await writeFile(
      join(tempDir, `${legacyRef}.json`),
      JSON.stringify(
        {
          data: {
            node_area_map: {
              J1: "DMA-1",
            },
          },
          projectId: "project-allowed",
        },
        null,
        2,
      ),
      "utf8",
    );

    const result = await resolver.getFullAuthorized(
      legacyRef,
      {
        actorKey: "actor-legacy-render",
        clientSessionId: "chat-legacy-render",
        projectId: "project-denied",
      },
      {
        expectedKind: RESULT_REFERENCE_KIND.renderJunctionsPayload,
      },
    );

    expect(result).toBeNull();
  });
});