Yonexus.Server/tests/connection-heartbeat-failures.test.ts

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

import {
  buildHeartbeat,
  buildHello,
  decodeBuiltin,
  encodeBuiltin,
  YONEXUS_PROTOCOL_VERSION
} from "../../Yonexus.Protocol/src/index.js";
import { createYonexusServerRuntime } from "../plugin/core/runtime.js";
import { createClientRecord, type ClientRecord } from "../plugin/core/persistence.js";
import type { YonexusServerStore } from "../plugin/core/store.js";
import type { ClientConnection, ServerTransport } from "../plugin/core/transport.js";

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 createMockStore(initialClients: ClientRecord[] = []): YonexusServerStore {
  const persisted = new Map(initialClients.map((record) => [record.identifier, record]));

  return {
    filePath: "/tmp/yonexus-server-connection-heartbeat-failures.json",
    load: vi.fn(async () => ({
      version: 1,
      persistedAt: 1_710_000_000,
      clients: new Map(persisted)
    })),
    save: vi.fn(async (clients: Iterable<ClientRecord>) => {
      persisted.clear();
      for (const client of clients) {
        persisted.set(client.identifier, client);
      }
    })
  };
}

function createMockTransport() {
  const sent: Array<{ connection: ClientConnection; message: string }> = [];
  const tempAssignments = new Map<ClientConnection["ws"], string>();
  const connections = new Map<string, ClientConnection>();

  const transport: ServerTransport = {
    isRunning: false,
    connections,
    start: vi.fn(),
    stop: vi.fn(),
    send: vi.fn(),
    sendToConnection: vi.fn((connection: ClientConnection, message: string) => {
      sent.push({ connection, message });
      return true;
    }),
    broadcast: vi.fn(),
    closeConnection: vi.fn(),
    promoteToAuthenticated: vi.fn((identifier: string, ws: ClientConnection["ws"]) => {
      if (!tempAssignments.has(ws)) {
        return false;
      }

      const existing = connections.get(identifier);
      if (existing) {
        existing.ws.close(1008, "Connection replaced by new authenticated session");
      }

      connections.set(identifier, {
        identifier,
        ws,
        connectedAt: 1_710_000_000,
        isAuthenticated: true
      });
      tempAssignments.delete(ws);
      return true;
    }),
    removeTempConnection: vi.fn((ws: ClientConnection["ws"]) => {
      tempAssignments.delete(ws);
    }),
    assignIdentifierToTemp: vi.fn((ws: ClientConnection["ws"], identifier: string) => {
      tempAssignments.set(ws, identifier);
    })
  };

  return { transport, sent };
}

describe("YNX-1105d: Connection & Heartbeat Failure Paths", () => {
  let now = 1_710_000_000;

  beforeEach(() => {
    now = 1_710_000_000;
    vi.useFakeTimers();
  });

  afterEach(() => {
    vi.useRealTimers();
  });

  it("CF-06: unauthenticated rule message closes connection", async () => {
    const record = createClientRecord("client-a");
    const store = createMockStore([record]);
    const { transport } = createMockTransport();

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

    await runtime.start();

    const connection = createConnection("client-a");
    runtime.state.registry.sessions.set("client-a", {
      identifier: "client-a",
      socket: connection.ws,
      isAuthenticated: false,
      connectedAt: now,
      lastActivityAt: now,
      publicKey: undefined
    });

    await runtime.handleMessage(connection, "chat::hello");

    expect(connection.ws.close).toHaveBeenCalledWith(1008, "Not authenticated");
  });

  it("HF-03: heartbeat before auth returns error", async () => {
    const record = createClientRecord("client-a");
    const store = createMockStore([record]);
    const { transport, sent } = createMockTransport();

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

    await runtime.start();

    const connection = createConnection("client-a");
    runtime.state.registry.sessions.set("client-a", {
      identifier: "client-a",
      socket: connection.ws,
      isAuthenticated: false,
      connectedAt: now,
      lastActivityAt: now,
      publicKey: undefined
    });

    await runtime.handleMessage(
      connection,
      encodeBuiltin(
        buildHeartbeat(
          { identifier: "client-a", status: "alive" },
          { requestId: "req-hb-early", timestamp: now }
        )
      )
    );

    const lastMessage = decodeBuiltin(sent.at(-1)!.message);
    expect(lastMessage.type).toBe("error");
    expect(lastMessage.payload).toMatchObject({
      code: "AUTH_FAILED"
    });
  });

  it("HF-04: heartbeat without session returns error", async () => {
    const record = createClientRecord("client-a");
    const store = createMockStore([record]);
    const { transport, sent } = createMockTransport();

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

    await runtime.start();

    const connection = createConnection("client-a");

    await runtime.handleMessage(
      connection,
      encodeBuiltin(
        buildHeartbeat(
          { identifier: "client-a", status: "alive" },
          { requestId: "req-hb-unauth", timestamp: now }
        )
      )
    );

    const lastMessage = decodeBuiltin(sent.at(-1)!.message);
    expect(lastMessage.type).toBe("error");
    expect(lastMessage.payload).toMatchObject({
      code: "AUTH_FAILED"
    });
  });

  it("CF-04: protocol version mismatch returns error and closes the connection", async () => {
    const record = createClientRecord("client-a");
    const store = createMockStore([record]);
    const { transport, sent } = createMockTransport();

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

    await runtime.start();

    const connection = createConnection();
    await runtime.handleMessage(
      connection,
      encodeBuiltin(
        buildHello(
          {
            identifier: "client-a",
            hasSecret: false,
            hasKeyPair: false,
            protocolVersion: `${YONEXUS_PROTOCOL_VERSION}-unsupported`
          },
          { requestId: "req-hello-version", timestamp: now }
        )
      )
    );

    const lastMessage = decodeBuiltin(sent.at(-1)!.message);
    expect(lastMessage.type).toBe("error");
    expect(lastMessage.payload).toMatchObject({
      code: "UNSUPPORTED_PROTOCOL_VERSION"
    });
    expect(connection.ws.close).toHaveBeenCalledWith(1002, "Unsupported protocol version");
  });

  it("CF-03: promoting a new authenticated connection replaces the old one", async () => {
    const previousSocket = createMockSocket();
    const replacementSocket = createMockSocket();
    const previousConnection = {
      identifier: "client-a",
      ws: previousSocket,
      connectedAt: now - 5,
      isAuthenticated: true
    } satisfies ClientConnection;

    const { transport } = createMockTransport();
    transport.connections.set("client-a", previousConnection);

    transport.assignIdentifierToTemp(replacementSocket, "client-a");
    const promoted = transport.promoteToAuthenticated("client-a", replacementSocket);

    expect(promoted).toBe(true);
    expect(previousSocket.close).toHaveBeenCalledWith(
      1008,
      "Connection replaced by new authenticated session"
    );

    const activeConnection = transport.connections.get("client-a");
    expect(activeConnection?.ws).toBe(replacementSocket);
    expect(activeConnection?.isAuthenticated).toBe(true);
  });
});