Yonexus.Server/tests/state-recovery.test.ts

import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { afterEach, describe, expect, it, vi } from "vitest";

import { createYonexusServerRuntime } from "../plugin/core/runtime.js";
import {
  createYonexusServerStore,
  loadServerStore,
  YonexusServerStoreCorruptionError
} from "../plugin/core/store.js";
import type { ClientConnection, ServerTransport } from "../plugin/core/transport.js";
import {
  buildHello,
  decodeBuiltin,
  encodeBuiltin,
  type BuiltinEnvelope,
  type PairRequestPayload,
  YONEXUS_PROTOCOL_VERSION
} from "../../Yonexus.Protocol/src/index.js";

const tempDirs: string[] = [];

afterEach(async () => {
  await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
});

async function createTempServerStorePath(): Promise<string> {
  const dir = await mkdtemp(join(tmpdir(), "yonexus-server-recovery-"));
  tempDirs.push(dir);
  return join(dir, "server-store.json");
}

function createMockSocket() {
  return {
    close: vi.fn()
  } as unknown as ClientConnection["ws"];
}

function createConnection(identifier: string | null = null): ClientConnection {
  return {
    identifier,
    ws: createMockSocket(),
    connectedAt: 1_710_000_000,
    isAuthenticated: false
  };
}

function createMockTransport() {
  const sentToConnection: Array<{ connection: ClientConnection; message: string }> = [];

  const transport: ServerTransport = {
    isRunning: false,
    connections: new Map(),
    start: vi.fn(async () => undefined),
    stop: vi.fn(async () => undefined),
    send: vi.fn(() => true),
    sendToConnection: vi.fn((connection: ClientConnection, message: string) => {
      sentToConnection.push({ connection, message });
      return true;
    }),
    broadcast: vi.fn(),
    closeConnection: vi.fn(() => true),
    promoteToAuthenticated: vi.fn(() => true),
    removeTempConnection: vi.fn(),
    assignIdentifierToTemp: vi.fn()
  };

  return {
    transport,
    sentToConnection
  };
}

describe("YNX-1105e: Server state recovery", () => {
  it("SR-01: preserves pending pairing across restart and reuses the same pairing code", async () => {
    const storePath = await createTempServerStorePath();
    const store = createYonexusServerStore(storePath);
    let now = 1_710_000_000;

    const firstTransport = createMockTransport();
    const firstRuntime = createYonexusServerRuntime({
      config: {
        followerIdentifiers: ["client-a"],
        notifyBotToken: "stub-token",
        adminUserId: "admin-user",
        listenHost: "127.0.0.1",
        listenPort: 8787
      },
      store,
      transport: firstTransport.transport,
      now: () => now
    });

    await firstRuntime.start();

    const firstConnection = createConnection();
    await firstRuntime.handleMessage(
      firstConnection,
      encodeBuiltin(
        buildHello(
          {
            identifier: "client-a",
            hasSecret: false,
            hasKeyPair: false,
            protocolVersion: YONEXUS_PROTOCOL_VERSION
          },
          { requestId: "req-hello-1", timestamp: now }
        )
      )
    );

    const initialRecord = firstRuntime.state.registry.clients.get("client-a");
    const initialPairingCode = initialRecord?.pairingCode;
    const initialExpiresAt = initialRecord?.pairingExpiresAt;

    expect(initialRecord?.pairingStatus).toBe("pending");
    expect(initialPairingCode).toBeTypeOf("string");
    expect(initialExpiresAt).toBeTypeOf("number");

    await firstRuntime.stop();

    const persistedRaw = JSON.parse(await readFile(storePath, "utf8")) as {
      clients: Array<{ identifier: string; pairingStatus: string; pairingCode?: string }>;
    };
    expect(
      persistedRaw.clients.find((client) => client.identifier === "client-a")
    ).toMatchObject({
      identifier: "client-a",
      pairingStatus: "pending",
      pairingCode: initialPairingCode
    });

    now += 30;

    const secondTransport = createMockTransport();
    const secondRuntime = createYonexusServerRuntime({
      config: {
        followerIdentifiers: ["client-a"],
        notifyBotToken: "stub-token",
        adminUserId: "admin-user",
        listenHost: "127.0.0.1",
        listenPort: 8787
      },
      store,
      transport: secondTransport.transport,
      now: () => now
    });

    await secondRuntime.start();

    const reloadedRecord = secondRuntime.state.registry.clients.get("client-a");
    expect(reloadedRecord?.pairingStatus).toBe("pending");
    expect(reloadedRecord?.pairingCode).toBe(initialPairingCode);
    expect(reloadedRecord?.pairingExpiresAt).toBe(initialExpiresAt);

    const secondConnection = createConnection();
    await secondRuntime.handleMessage(
      secondConnection,
      encodeBuiltin(
        buildHello(
          {
            identifier: "client-a",
            hasSecret: false,
            hasKeyPair: false,
            protocolVersion: YONEXUS_PROTOCOL_VERSION
          },
          { requestId: "req-hello-2", timestamp: now }
        )
      )
    );

    const helloAck = decodeBuiltin(secondTransport.sentToConnection[0].message);
    expect(helloAck.type).toBe("hello_ack");
    expect(helloAck.payload).toMatchObject({
      identifier: "client-a",
      nextAction: "waiting_pair_confirm"
    });

    const pairRequest = decodeBuiltin(secondTransport.sentToConnection[1].message) as BuiltinEnvelope<
      "pair_request",
      PairRequestPayload
    >;
    expect(pairRequest.type).toBe("pair_request");
    expect(pairRequest.payload).toMatchObject({
      identifier: "client-a",
      adminNotification: "sent",
      codeDelivery: "out_of_band"
    });
    expect(pairRequest.payload?.expiresAt).toBe(initialExpiresAt);

    await secondRuntime.stop();
  });

  it("SR-02: restart drops in-memory active sessions and requires reconnect", async () => {
    const storePath = await createTempServerStorePath();
    const store = createYonexusServerStore(storePath);
    const now = 1_710_000_000;

    const firstTransport = createMockTransport();
    const firstRuntime = createYonexusServerRuntime({
      config: {
        followerIdentifiers: ["client-a"],
        notifyBotToken: "stub-token",
        adminUserId: "admin-user",
        listenHost: "127.0.0.1",
        listenPort: 8787
      },
      store,
      transport: firstTransport.transport,
      now: () => now
    });

    await firstRuntime.start();
    const record = firstRuntime.state.registry.clients.get("client-a");
    expect(record).toBeDefined();

    record!.pairingStatus = "paired";
    record!.publicKey = "test-public-key";
    record!.secret = "test-secret";
    record!.status = "online";
    record!.lastAuthenticatedAt = now;
    record!.lastHeartbeatAt = now;
    record!.updatedAt = now;

    firstRuntime.state.registry.sessions.set("client-a", {
      identifier: "client-a",
      socket: createMockSocket(),
      isAuthenticated: true,
      connectedAt: now,
      lastActivityAt: now,
      publicKey: "test-public-key"
    });

    await firstRuntime.stop();

    const secondTransport = createMockTransport();
    const secondRuntime = createYonexusServerRuntime({
      config: {
        followerIdentifiers: ["client-a"],
        notifyBotToken: "stub-token",
        adminUserId: "admin-user",
        listenHost: "127.0.0.1",
        listenPort: 8787
      },
      store,
      transport: secondTransport.transport,
      now: () => now + 5
    });

    await secondRuntime.start();

    const reloadedRecord = secondRuntime.state.registry.clients.get("client-a");
    expect(reloadedRecord).toMatchObject({
      identifier: "client-a",
      pairingStatus: "paired",
      secret: "test-secret",
      publicKey: "test-public-key",
      status: "online",
      lastAuthenticatedAt: now,
      lastHeartbeatAt: now
    });
    expect(secondRuntime.state.registry.sessions.size).toBe(0);

    const reconnectConnection = createConnection();
    await secondRuntime.handleMessage(
      reconnectConnection,
      encodeBuiltin(
        buildHello(
          {
            identifier: "client-a",
            hasSecret: true,
            hasKeyPair: true,
            publicKey: "test-public-key",
            protocolVersion: YONEXUS_PROTOCOL_VERSION
          },
          { requestId: "req-hello-reconnect", timestamp: now + 5 }
        )
      )
    );

    const helloAck = decodeBuiltin(secondTransport.sentToConnection[0].message);
    expect(helloAck.type).toBe("hello_ack");
    expect(helloAck.payload).toMatchObject({
      identifier: "client-a",
      nextAction: "auth_required"
    });

    await secondRuntime.stop();
  });

  it("SR-05: corrupted server store raises YonexusServerStoreCorruptionError", async () => {
    const storePath = await createTempServerStorePath();
    await writeFile(storePath, '{"version":1,"clients":"oops"}\n', "utf8");

    await expect(loadServerStore(storePath)).rejects.toBeInstanceOf(YonexusServerStoreCorruptionError);
    await expect(loadServerStore(storePath)).rejects.toThrow("invalid clients array");
  });
});