Yonexus.Client/tests/runtime-flow.test.ts

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

import {
  buildAuthFailed,
  buildAuthSuccess,
  buildHelloAck,
  buildPairFailed,
  buildPairRequest,
  buildPairSuccess,
  buildRePairRequired,
  decodeBuiltin,
  encodeBuiltin,
  type HelloEnvelopePayloadMap,
  type PairConfirmPayload,
  type TypedBuiltinEnvelope
} from "../../Yonexus.Protocol/src/index.js";
import { createYonexusClientRuntime } from "../plugin/core/runtime.js";
import type { YonexusClientState, YonexusClientStateStore } from "../plugin/core/state.js";
import type { ClientConnectionState, ClientTransport } from "../plugin/core/transport.js";

type SavedState = YonexusClientState;

function createInitialState(): YonexusClientState {
  return {
    identifier: "client-a",
    updatedAt: 1_710_000_000
  };
}

function createMockStateStore(initialState: YonexusClientState = createInitialState()) {
  let state = { ...initialState };
  const saved: SavedState[] = [];

  const store: YonexusClientStateStore = {
    filePath: "/tmp/yonexus-client-test.json",
    load: vi.fn(async () => ({ ...state })),
    save: vi.fn(async (next) => {
      state = { ...next };
      saved.push({ ...next });
    })
  };

  return {
    store,
    saved,
    getState: () => ({ ...state })
  };
}

function createMockTransport() {
  let currentState: ClientConnectionState = "idle";
  const sent: string[] = [];

  const transport: ClientTransport = {
    get state() {
      return currentState;
    },
    get isConnected() {
      return currentState !== "idle" && currentState !== "disconnected" && currentState !== "error";
    },
    get isAuthenticated() {
      return currentState === "authenticated";
    },
    connect: vi.fn(async () => {
      currentState = "connected";
    }),
    disconnect: vi.fn(() => {
      currentState = "disconnected";
    }),
    send: vi.fn((message: string) => {
      sent.push(message);
      return true;
    }),
    markAuthenticated: vi.fn(() => {
      currentState = "authenticated";
    }),
    markAuthenticating: vi.fn(() => {
      currentState = "authenticating";
    })
  };

  return {
    transport,
    sent,
    setState: (state: ClientConnectionState) => {
      currentState = state;
    }
  };
}

describe("Yonexus.Client runtime flow", () => {
  it("starts by loading state, ensuring keypair, and sending hello on connect", async () => {
    const storeState = createMockStateStore();
    const transportState = createMockTransport();
    const runtime = createYonexusClientRuntime({
      config: {
        mainHost: "ws://localhost:8787",
        identifier: "client-a",
        notifyBotToken: "stub-token",
        adminUserId: "admin-user"
      },
      transport: transportState.transport,
      stateStore: storeState.store,
      now: () => 1_710_000_000
    });

    await runtime.start();
    runtime.handleTransportStateChange("connected");

    expect(transportState.transport.connect).toHaveBeenCalled();
    expect(storeState.saved.length).toBeGreaterThan(0);
    expect(runtime.state.clientState.publicKey).toBeTypeOf("string");
    expect(runtime.state.phase).toBe("awaiting_hello_ack");

    const hello = decodeBuiltin(transportState.sent[0]);
    expect(hello.type).toBe("hello");
    expect(hello.payload).toMatchObject({
      identifier: "client-a",
      hasSecret: false,
      hasKeyPair: true
    });
  });

  it("handles pair request, submits code, stores secret, and authenticates", async () => {
    let now = 1_710_000_000;
    const storeState = createMockStateStore();
    const transportState = createMockTransport();
    const runtime = createYonexusClientRuntime({
      config: {
        mainHost: "ws://localhost:8787",
        identifier: "client-a",
        notifyBotToken: "stub-token",
        adminUserId: "admin-user"
      },
      transport: transportState.transport,
      stateStore: storeState.store,
      now: () => now
    });

    await runtime.start();
    runtime.handleTransportStateChange("connected");
    await runtime.handleMessage(
      encodeBuiltin(
        buildHelloAck(
          {
            identifier: "client-a",
            nextAction: "pair_required"
          },
          { requestId: "req-hello", timestamp: now }
        )
      )
    );

    expect(runtime.state.phase).toBe("pair_required");

    await runtime.handleMessage(
      encodeBuiltin(
        buildPairRequest(
          {
            identifier: "client-a",
            expiresAt: now + 300,
            ttlSeconds: 300,
            adminNotification: "sent",
            codeDelivery: "out_of_band"
          },
          { requestId: "req-pair", timestamp: now }
        )
      )
    );

    expect(runtime.state.phase).toBe("waiting_pair_confirm");
    expect(runtime.state.pendingPairing).toMatchObject({
      ttlSeconds: 300,
      adminNotification: "sent"
    });

    expect(runtime.submitPairingCode("PAIR-CODE-123", "req-pair-confirm")).toBe(true);
    const pairConfirm = decodeBuiltin(transportState.sent.at(-1)!);
    expect(pairConfirm.type).toBe("pair_confirm");
    expect((pairConfirm.payload as PairConfirmPayload).pairingCode).toBe("PAIR-CODE-123");

    now += 1;
    await runtime.handleMessage(
      encodeBuiltin(
        buildPairSuccess(
          {
            identifier: "client-a",
            secret: "issued-secret",
            pairedAt: now
          },
          { requestId: "req-pair-confirm", timestamp: now }
        )
      )
    );

    expect(runtime.state.clientState.secret).toBe("issued-secret");
    expect(runtime.state.phase).toBe("auth_required");
    expect(transportState.transport.markAuthenticating).toHaveBeenCalled();

    const authRequest = decodeBuiltin(transportState.sent.at(-1)!);
    expect(authRequest.type).toBe("auth_request");
    expect(authRequest.payload).toMatchObject({ identifier: "client-a" });

    now += 1;
    await runtime.handleMessage(
      encodeBuiltin(
        buildAuthSuccess(
          {
            identifier: "client-a",
            authenticatedAt: now,
            status: "online"
          },
          { requestId: "req-auth", timestamp: now }
        )
      )
    );

    expect(runtime.state.phase).toBe("authenticated");
    expect(transportState.transport.markAuthenticated).toHaveBeenCalled();
    expect(runtime.state.clientState.authenticatedAt).toBe(now);
  });

  it("resets trust state on re-pair-required auth failures", async () => {
    let now = 1_710_000_000;
    const storeState = createMockStateStore({
      identifier: "client-a",
      publicKey: "pubkey",
      privateKey: "privkey",
      secret: "old-secret",
      pairedAt: now - 10,
      authenticatedAt: now - 5,
      updatedAt: now - 5
    });
    const transportState = createMockTransport();
    transportState.setState("connected");
    const runtime = createYonexusClientRuntime({
      config: {
        mainHost: "ws://localhost:8787",
        identifier: "client-a",
        notifyBotToken: "stub-token",
        adminUserId: "admin-user"
      },
      transport: transportState.transport,
      stateStore: storeState.store,
      now: () => now
    });

    await runtime.start();
    await runtime.handleMessage(
      encodeBuiltin(
        buildAuthFailed(
          {
            identifier: "client-a",
            reason: "nonce_collision"
          },
          { requestId: "req-auth", timestamp: now }
        )
      )
    );

    expect(runtime.state.phase).toBe("pair_required");
    expect(runtime.state.lastPairingFailure).toBe("nonce_collision");
    expect(runtime.state.clientState.secret).toBeUndefined();

    now += 1;
    await runtime.handleMessage(
      encodeBuiltin(
        buildRePairRequired(
          {
            identifier: "client-a",
            reason: "rate_limited"
          },
          { requestId: "req-repair", timestamp: now }
        )
      )
    );

    expect(runtime.state.phase).toBe("pair_required");
    expect(runtime.state.lastPairingFailure).toBe("re_pair_required");
  });

  it("SR-03: restarts with stored credentials and resumes at auth flow without re-pairing", async () => {
    const now = 1_710_000_000;
    const { generateKeyPair } = await import("../plugin/crypto/keypair.js");
    const keyPair = await generateKeyPair();
    const storeState = createMockStateStore({
      identifier: "client-a",
      publicKey: keyPair.publicKey,
      privateKey: keyPair.privateKey,
      secret: "stored-secret",
      pairedAt: now - 20,
      authenticatedAt: now - 10,
      updatedAt: now - 10
    });
    const transportState = createMockTransport();
    const runtime = createYonexusClientRuntime({
      config: {
        mainHost: "ws://localhost:8787",
        identifier: "client-a",
        notifyBotToken: "stub-token",
        adminUserId: "admin-user"
      },
      transport: transportState.transport,
      stateStore: storeState.store,
      now: () => now
    });

    await runtime.start();
    runtime.handleTransportStateChange("connected");

    const hello = decodeBuiltin(transportState.sent[0]);
    expect(hello.type).toBe("hello");
    expect(hello.payload).toMatchObject({
      identifier: "client-a",
      hasSecret: true,
      hasKeyPair: true,
      publicKey: keyPair.publicKey
    });

    await runtime.handleMessage(
      encodeBuiltin(
        buildHelloAck(
          {
            identifier: "client-a",
            nextAction: "auth_required"
          },
          { requestId: "req-restart-hello", timestamp: now }
        )
      )
    );

    expect(runtime.state.phase).toBe("auth_required");

    const authRequest = decodeBuiltin(transportState.sent.at(-1)!);
    expect(authRequest.type).toBe("auth_request");
    expect(authRequest.payload).toMatchObject({ identifier: "client-a" });
  });

  it("sends heartbeat only when authenticated and connected", async () => {
    const storeState = createMockStateStore({
      identifier: "client-a",
      publicKey: "pubkey",
      privateKey: "privkey",
      secret: "secret",
      pairedAt: 1_709_999_990,
      authenticatedAt: 1_709_999_995,
      updatedAt: 1_709_999_995
    });
    const transportState = createMockTransport();
    transportState.setState("authenticated");
    const runtime = createYonexusClientRuntime({
      config: {
        mainHost: "ws://localhost:8787",
        identifier: "client-a",
        notifyBotToken: "stub-token",
        adminUserId: "admin-user"
      },
      transport: transportState.transport,
      stateStore: storeState.store,
      now: () => 1_710_000_000
    });

    await runtime.start();

    await runtime.handleMessage("heartbeat_tick");
    expect(transportState.sent).toHaveLength(0);

    await runtime.handleMessage(
      encodeBuiltin(
        buildAuthSuccess(
          {
            identifier: "client-a",
            authenticatedAt: 1_710_000_000,
            status: "online"
          },
          { timestamp: 1_710_000_000 }
        )
      )
    );

    await runtime.handleMessage("heartbeat_tick");
    const heartbeat = decodeBuiltin(transportState.sent.at(-1)!);
    expect(heartbeat.type).toBe("heartbeat");
    expect(heartbeat.payload).toMatchObject({ identifier: "client-a", status: "alive" });
  });

  it("tracks pairing failures without wiping pending session for retryable reasons", async () => {
    const storeState = createMockStateStore();
    const transportState = createMockTransport();
    const runtime = createYonexusClientRuntime({
      config: {
        mainHost: "ws://localhost:8787",
        identifier: "client-a",
        notifyBotToken: "stub-token",
        adminUserId: "admin-user"
      },
      transport: transportState.transport,
      stateStore: storeState.store,
      now: () => 1_710_000_000
    });

    await runtime.start();
    runtime.handleTransportStateChange("connected");
    await runtime.handleMessage(
      encodeBuiltin(
        buildPairRequest(
          {
            identifier: "client-a",
            expiresAt: 1_710_000_300,
            ttlSeconds: 300,
            adminNotification: "sent",
            codeDelivery: "out_of_band"
          },
          { timestamp: 1_710_000_000 }
        )
      )
    );

    await runtime.handleMessage(
      encodeBuiltin(
        buildPairFailed(
          {
            identifier: "client-a",
            reason: "invalid_code"
          },
          { timestamp: 1_710_000_001 }
        )
      )
    );

    expect(runtime.state.phase).toBe("waiting_pair_confirm");
    expect(runtime.state.pendingPairing).toBeDefined();
    expect(runtime.state.lastPairingFailure).toBe("invalid_code");
  });
});