polyguard-rl/app/ui/frontend/src/lib/api.ts

import { API_BASE, ENV_BASE } from "./constants";
import type {
  EnvCatalog,
  EnvWsMessage,
  EnvObservation,
  ModelStatus,
  ResetEnvOptions,
  StepCandidatePayload,
  StepResponse,
} from "./types";

async function fetchJson<T>(path: string, init?: RequestInit): Promise<T> {
  const res = await fetch(`${API_BASE}${path}`, init);
  if (!res.ok) {
    const body = await res.text();
    throw new Error(`API ${path} failed (${res.status}): ${body.slice(0, 240)}`);
  }
  return (await res.json()) as T;
}

let envSocket: WebSocket | null = null;
const pendingEnvMessages: Array<{
  resolve: (value: unknown) => void;
  reject: (reason: unknown) => void;
}> = [];

function envWsUrl(): string {
  const base = ENV_BASE.replace(/\/$/, "");
  return `${base.replace(/^http/, "ws")}/ws`;
}

async function ensureEnvSocket(): Promise<WebSocket> {
  if (envSocket?.readyState === WebSocket.OPEN) return envSocket;
  if (envSocket?.readyState === WebSocket.CONNECTING) {
    await new Promise((resolve) => setTimeout(resolve, 80));
    return ensureEnvSocket();
  }

  const socket = new WebSocket(envWsUrl());
  envSocket = socket;

  socket.onmessage = (event) => {
    const pending = pendingEnvMessages.shift();
    if (!pending) return;
    try {
      const message = JSON.parse(event.data as string) as EnvWsMessage;
      if (message.type === "error") {
        const data = message.data;
        const messageText =
          data && typeof data === "object" && "message" in data
            ? String((data as Record<string, unknown>).message)
            : "Env service returned an error";
        pending.reject(new Error(messageText));
        return;
      }
      pending.resolve(message.data);
    } catch (err) {
      pending.reject(err);
    }
  };

  socket.onerror = () => {
    const pending = pendingEnvMessages.shift();
    if (pending) pending.reject(new Error(`Unable to connect to env service at ${envWsUrl()}`));
  };

  socket.onclose = () => {
    envSocket = null;
  };

  await new Promise<void>((resolve, reject) => {
    const timeout = window.setTimeout(() => reject(new Error(`Env service timeout at ${envWsUrl()}`)), 2500);
    socket.onopen = () => {
      window.clearTimeout(timeout);
      resolve();
    };
  });

  return socket;
}

export async function envWsSend<T>(type: string, data: unknown): Promise<T> {
  const socket = await ensureEnvSocket();
  return new Promise<T>((resolve, reject) => {
    pendingEnvMessages.push({
      resolve: (value) => resolve(value as T),
      reject,
    });
    socket.send(JSON.stringify({ type, data }));
  });
}

export function closeEnvSocket(): void {
  try {
    envSocket?.close();
  } catch {
    // Ignore close errors during route teardown.
  } finally {
    envSocket = null;
    pendingEnvMessages.splice(0);
  }
}

export async function fetchCatalog(): Promise<EnvCatalog> {
  return fetchJson<EnvCatalog>("/env/catalog");
}

export async function resetEnv(options: ResetEnvOptions = {}): Promise<EnvObservation> {
  return fetchJson<EnvObservation>("/env/reset", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(options),
  });
}

export async function orchestrateStep(): Promise<Record<string, unknown>> {
  return fetchJson<Record<string, unknown>>("/agents/orchestrate", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({}),
  });
}

export async function stepCandidate(payload: StepCandidatePayload): Promise<StepResponse> {
  return fetchJson<StepResponse>("/env/step_candidate", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(payload),
  });
}

export async function fetchTrace(): Promise<Array<Record<string, unknown>>> {
  return fetchJson<Array<Record<string, unknown>>>("/env/trace");
}

export async function fetchTrainingMetrics(): Promise<Record<string, unknown>> {
  return fetchJson<Record<string, unknown>>("/metrics/training");
}

export async function fetchBaselines(): Promise<Record<string, unknown>> {
  return fetchJson<Record<string, unknown>>("/eval/run_baselines", { method: "POST" });
}

export async function fetchRewardBreakdown(): Promise<Record<string, unknown>> {
  return fetchJson<Record<string, unknown>>("/env/reward_breakdown");
}

export async function fetchModelStatus(): Promise<ModelStatus> {
  return fetchJson<ModelStatus>("/policy/model_status");
}

export async function fetchLegalActions(): Promise<Array<Record<string, unknown>>> {
  return fetchJson<Array<Record<string, unknown>>>("/env/legal_actions");
}

export async function fetchUncertainty(): Promise<Record<string, unknown>> {
  return fetchJson<Record<string, unknown>>("/env/uncertainty");
}

export async function fetchDosingEval(): Promise<Record<string, unknown>> {
  return fetchJson<Record<string, unknown>>("/eval/run_dosing", { method: "POST" });
}