Yonexus.Client/plugin/core/runtime.ts

import {
  YONEXUS_PROTOCOL_VERSION,
  buildAuthRequest,
  buildHeartbeat,
  buildHello,
  buildPairConfirm,
  CodecError,
  createAuthRequestSigningInput,
  decodeBuiltin,
  encodeBuiltin,
  encodeRuleMessage,
  isBuiltinMessage,
  type AuthFailedPayload,
  type BuiltinPayloadMap,
  type HelloAckPayload,
  type PairFailedPayload,
  type PairRequestPayload,
  type PairSuccessPayload,
  type TypedBuiltinEnvelope
} from "../../../Yonexus.Protocol/src/index.js";
import type { YonexusClientConfig } from "./config.js";
import {
  ensureClientKeyPair,
  hasClientKeyPair,
  hasClientSecret,
  createInitialClientState,
  YonexusClientState,
  YonexusClientStateStore
} from "./state.js";
import { generateNonce, signMessage } from "../crypto/keypair.js";
import type { ClientConnectionState, ClientTransport } from "./transport.js";
import type { ClientRuleRegistry } from "./rules.js";

export type YonexusClientPhase =
  | "idle"
  | "starting"
  | "awaiting_hello_ack"
  | "pair_required"
  | "waiting_pair_confirm"
  | "auth_required"
  | "authenticated"
  | "stopped";

export interface YonexusClientRuntimeOptions {
  config: YonexusClientConfig;
  transport: ClientTransport;
  stateStore: YonexusClientStateStore;
  ruleRegistry?: ClientRuleRegistry;
  onAuthenticated?: () => void;
  now?: () => number;
}

export interface YonexusClientRuntimeState {
  readonly phase: YonexusClientPhase;
  readonly transportState: ClientConnectionState;
  readonly clientState: YonexusClientState;
  readonly pendingPairing?: {
    expiresAt: number;
    ttlSeconds: number;
    adminNotification: "sent" | "failed";
  };
  readonly lastPairingFailure?: string;
}

export class YonexusClientRuntime {
  private readonly options: YonexusClientRuntimeOptions;
  private readonly now: () => number;
  private clientState: YonexusClientState;
  private phase: YonexusClientPhase = "idle";
  private pendingPairing?: {
    expiresAt: number;
    ttlSeconds: number;
    adminNotification: "sent" | "failed";
  };
  private lastPairingFailure?: string;

  constructor(options: YonexusClientRuntimeOptions) {
    this.options = options;
    this.now = options.now ?? (() => Math.floor(Date.now() / 1000));
    this.clientState = createInitialClientState(options.config.identifier);
  }

  get state(): YonexusClientRuntimeState {
    return {
      phase: this.phase,
      transportState: this.options.transport.state,
      clientState: this.clientState,
      pendingPairing: this.pendingPairing,
      lastPairingFailure: this.lastPairingFailure
    };
  }

  async start(): Promise<void> {
    if (this.phase !== "idle" && this.phase !== "stopped") {
      return;
    }

    this.phase = "starting";
    
    // Load existing state and ensure key pair exists
    let state = await this.options.stateStore.load(this.options.config.identifier);
    const keyResult = await ensureClientKeyPair(state, this.options.stateStore);
    this.clientState = keyResult.state;
    
    await this.options.transport.connect();
  }

  async stop(): Promise<void> {
    await this.options.stateStore.save({
      ...this.clientState,
      updatedAt: this.now()
    });
    this.options.transport.disconnect();
    this.phase = "stopped";
  }

  async handleMessage(raw: string): Promise<void> {
    if (raw === "heartbeat_tick") {
      await this.handleHeartbeatTick();
      return;
    }

    if (!isBuiltinMessage(raw)) {
      this.options.ruleRegistry?.dispatch(raw);
      return;
    }

    let envelope: TypedBuiltinEnvelope<keyof BuiltinPayloadMap>;
    try {
      envelope = decodeBuiltin(raw) as TypedBuiltinEnvelope<keyof BuiltinPayloadMap>;
    } catch (error) {
      if (error instanceof CodecError) {
        this.lastPairingFailure = error.message;
      }
      return;
    }
    if (envelope.type === "hello_ack") {
      this.handleHelloAck(envelope as TypedBuiltinEnvelope<"hello_ack">);
      return;
    }

    if (envelope.type === "pair_request") {
      this.handlePairRequest(envelope as TypedBuiltinEnvelope<"pair_request">);
      return;
    }

    if (envelope.type === "pair_success") {
      await this.handlePairSuccess(envelope as TypedBuiltinEnvelope<"pair_success">);
      return;
    }

    if (envelope.type === "pair_failed") {
      this.handlePairFailed(envelope as TypedBuiltinEnvelope<"pair_failed">);
      return;
    }

    if (envelope.type === "auth_success") {
      this.options.transport.markAuthenticated();
      this.clientState = {
        ...this.clientState,
        authenticatedAt: this.now(),
        updatedAt: this.now()
      };
      await this.options.stateStore.save(this.clientState);
      this.phase = "authenticated";
      this.options.onAuthenticated?.();
      return;
    }

    if (envelope.type === "auth_failed") {
      await this.handleAuthFailed(envelope as TypedBuiltinEnvelope<"auth_failed">);
      return;
    }

    if (envelope.type === "re_pair_required") {
      await this.handleRePairRequired();
      return;
    }

    this.lastPairingFailure = `unsupported_builtin:${String(envelope.type)}`;
  }

  handleTransportStateChange(state: ClientConnectionState): void {
    if (state === "connected") {
      // Reload state from disk before hello so that any secret written by an
      // external process (e.g. a pairing script) is picked up on reconnect.
      this.options.stateStore.load(this.options.config.identifier).then((fresh) => {
        if (fresh) {
          this.clientState = { ...this.clientState, ...fresh };
        }
        this.sendHello();
      }).catch(() => {
        // If reload fails, proceed with in-memory state
        this.sendHello();
      });
    }

    if (state === "disconnected") {
      this.phase = hasClientSecret(this.clientState) ? "auth_required" : "idle";
      this.pendingPairing = undefined;
    }
  }

  private sendHello(): void {
    this.phase = "awaiting_hello_ack";
    this.options.transport.send(
      encodeBuiltin(
        buildHello(
          {
            identifier: this.options.config.identifier,
            hasSecret: hasClientSecret(this.clientState),
            hasKeyPair: hasClientKeyPair(this.clientState),
            publicKey: this.clientState.publicKey,
            protocolVersion: YONEXUS_PROTOCOL_VERSION
          },
          { timestamp: this.now() }
        )
      )
    );
  }

  submitPairingCode(pairingCode: string, requestId?: string): boolean {
    const normalizedCode = pairingCode.trim();
    if (!normalizedCode || !this.options.transport.isConnected) {
      return false;
    }

    this.lastPairingFailure = undefined;
    return this.options.transport.send(
      encodeBuiltin(
        buildPairConfirm(
          {
            identifier: this.options.config.identifier,
            pairingCode: normalizedCode
          },
          { requestId, timestamp: this.now() }
        )
      )
    );
  }

  private handleHelloAck(envelope: TypedBuiltinEnvelope<"hello_ack">): void {
    const payload = envelope.payload as HelloAckPayload | undefined;
    if (!payload) {
      return;
    }

    switch (payload.nextAction) {
      case "pair_required":
        this.phase = "pair_required";
        break;
      case "waiting_pair_confirm":
        this.phase = "waiting_pair_confirm";
        break;
      case "auth_required":
        this.phase = "auth_required";
        void this.sendAuthRequest();
        break;
      default:
        this.phase = "idle";
        break;
    }
  }

  private handlePairRequest(envelope: TypedBuiltinEnvelope<"pair_request">): void {
    const payload = envelope.payload as PairRequestPayload | undefined;
    if (!payload) {
      return;
    }

    this.pendingPairing = {
      expiresAt: payload.expiresAt,
      ttlSeconds: payload.ttlSeconds,
      adminNotification: payload.adminNotification
    };
    this.lastPairingFailure = undefined;
    // Always wait for the pairing code regardless of notification status.
    // When adminNotification is "failed", the admin can retrieve the code
    // via the server CLI command and deliver it through an alternate channel.
    this.phase = "waiting_pair_confirm";
  }

  private async handlePairSuccess(envelope: TypedBuiltinEnvelope<"pair_success">): Promise<void> {
    const payload = envelope.payload as PairSuccessPayload | undefined;
    if (!payload) {
      return;
    }

    this.clientState = {
      ...this.clientState,
      secret: payload.secret,
      pairedAt: payload.pairedAt,
      updatedAt: this.now()
    };
    await this.options.stateStore.save(this.clientState);
    this.pendingPairing = undefined;
    this.lastPairingFailure = undefined;
    this.phase = "auth_required";
    await this.sendAuthRequest();
  }

  private handlePairFailed(envelope: TypedBuiltinEnvelope<"pair_failed">): void {
    const payload = envelope.payload as PairFailedPayload | undefined;
    if (!payload) {
      return;
    }

    this.lastPairingFailure = payload.reason;
    if (payload.reason === "expired" || payload.reason === "admin_notification_failed") {
      this.pendingPairing = undefined;
      this.phase = "pair_required";
      return;
    }

    this.phase = "waiting_pair_confirm";
  }

  private async handleAuthFailed(
    envelope: TypedBuiltinEnvelope<"auth_failed">
  ): Promise<void> {
    const payload = envelope.payload as AuthFailedPayload | undefined;
    if (!payload) {
      return;
    }

    if (
      payload.reason === "re_pair_required" ||
      payload.reason === "nonce_collision" ||
      payload.reason === "rate_limited"
    ) {
      this.lastPairingFailure = payload.reason;
      await this.resetTrustState();
      return;
    }

    this.lastPairingFailure = payload.reason;
    // If the server lost our session (race condition), re-announce via hello
    // so the server creates a new session and we can retry auth.
    if (payload.reason === "not_paired" && hasClientSecret(this.clientState)) {
      this.sendHello();
      return;
    }
    this.phase = "auth_required";
  }

  private async sendAuthRequest(): Promise<void> {
    if (!this.options.transport.isConnected) {
      return;
    }

    if (!this.clientState.secret || !this.clientState.privateKey) {
      this.phase = "pair_required";
      return;
    }

    const proofTimestamp = this.now();
    const nonce = generateNonce();
    const signature = await signMessage(
      this.clientState.privateKey,
      createAuthRequestSigningInput({
        secret: this.clientState.secret,
        nonce,
        proofTimestamp
      })
    );

    this.options.transport.markAuthenticating();
    this.options.transport.send(
      encodeBuiltin(
        buildAuthRequest(
          {
            identifier: this.options.config.identifier,
            nonce,
            proofTimestamp,
            signature,
            publicKey: this.clientState.publicKey
          },
          { requestId: `auth_${proofTimestamp}_${nonce}`, timestamp: proofTimestamp }
        )
      )
    );
  }

  private async handleHeartbeatTick(): Promise<void> {
    if (this.phase !== "authenticated" || !this.options.transport.isConnected) {
      return;
    }

    this.options.transport.send(
      encodeBuiltin(
        buildHeartbeat(
          {
            identifier: this.options.config.identifier,
            status: "alive"
          },
          { timestamp: this.now() }
        )
      )
    );
  }

  private async handleRePairRequired(): Promise<void> {
    this.pendingPairing = undefined;
    this.lastPairingFailure = "re_pair_required";
    await this.resetTrustState();
  }

  private async resetTrustState(): Promise<void> {
    this.clientState = {
      ...this.clientState,
      secret: undefined,
      pairedAt: undefined,
      authenticatedAt: undefined,
      updatedAt: this.now()
    };

    await this.options.stateStore.save(this.clientState);
    this.phase = "pair_required";
  }

  /**
   * Send a rule message to the server.
   * Message must already conform to `${rule_identifier}::${message_content}`.
   *
   * @param message - The complete rule message with identifier and content
   * @returns True if message was sent, false if not connected or not authenticated
   */
  sendMessageToServer(message: string): boolean {
    if (!this.options.transport.isConnected || !this.options.transport.isAuthenticated) {
      return false;
    }

    // Validate the message is a properly formatted rule message
    try {
      if (message.startsWith("builtin::")) {
        return false;
      }
      const delimiterIndex = message.indexOf("::");
      if (delimiterIndex === -1) {
        return false;
      }
      const ruleIdentifier = message.slice(0, delimiterIndex);
      const content = message.slice(delimiterIndex + 2);
      encodeRuleMessage(ruleIdentifier, content);
    } catch {
      return false;
    }

    return this.options.transport.send(message);
  }

  /**
   * Send a rule message to the server using separate rule identifier and content.
   *
   * @param ruleIdentifier - The rule identifier (alphanumeric with underscores/hyphens)
   * @param content - The message content
   * @returns True if message was sent, false if not connected/authenticated or invalid format
   */
  sendRuleMessage(ruleIdentifier: string, content: string): boolean {
    if (!this.options.transport.isConnected || !this.options.transport.isAuthenticated) {
      return false;
    }

    try {
      const encoded = encodeRuleMessage(ruleIdentifier, content);
      return this.options.transport.send(encoded);
    } catch {
      return false;
    }
  }
}

export function createYonexusClientRuntime(
  options: YonexusClientRuntimeOptions
): YonexusClientRuntime {
  return new YonexusClientRuntime(options);
}