next-tjwater-drainage-frontend/features/agent/ui-envelope/validator.ts

import { normalizeSafeChartData, parseChartSpec } from "../chart-data";
import { toTrustedMapAction } from "./map-actions";
import type { UIEnvelope, UIEnvelopePayload, UIRegistry, UISurface } from "./types";

const surfaces = new Set<UISurface>(["chat_inline", "side_panel", "canvas", "map_overlay"]);
const MAX_ID_LENGTH = 128;
const LOCAL_COMPONENT_SURFACES = new Map<string, UISurface[]>([
  ["HistoryPanel", ["side_panel", "canvas"]]
]);
const LOCAL_ACTION_SURFACES = new Map<string, UISurface[]>([
  ["locate_features", ["map_overlay"]],
  ["zoom_to_map", ["map_overlay"]],
  ["render_junctions", ["map_overlay"]],
  ["apply_layer_style", ["map_overlay"]]
]);

export function parseUiEnvelopePayload(value: unknown): UIEnvelopePayload | null {
  if (!isRecord(value)) {
    return null;
  }

  const envelope = parseUiEnvelope(value.envelope);
  if (!envelope) {
    return null;
  }

  const sessionId = boundedString(value.session_id, MAX_ID_LENGTH);
  const envelopeId = boundedString(value.envelope_id, MAX_ID_LENGTH);
  if (!sessionId || !envelopeId || !isFiniteTimestamp(value.created_at)) {
    return null;
  }

  return {
    session_id: sessionId,
    envelope_id: envelopeId,
    created_at: value.created_at,
    envelope
  };
}

export function parseUiEnvelope(value: unknown): UIEnvelope | null {
  if (!isRecord(value) || value.schemaVersion !== "agent-ui@1") {
    return null;
  }

  if (value.kind === "registered_component") {
    if (typeof value.component !== "string" || !isSurface(value.surface) || !isValidComponentProps(value.component, value.props)) {
      return null;
    }

    return {
      kind: "registered_component",
      schemaVersion: "agent-ui@1",
      component: value.component,
      surface: value.surface,
      props: value.props,
      data: value.data,
      fallbackText: optionalString(value.fallbackText)
    };
  }

  if (value.kind === "chart") {
    if (value.grammar !== "echarts-safe-subset" || !isChartSurface(value.surface) || !isValidChart(value.spec, value.data)) {
      return null;
    }

    return {
      kind: "chart",
      schemaVersion: "agent-ui@1",
      grammar: "echarts-safe-subset",
      surface: value.surface,
      spec: value.spec,
      data: value.data,
      fallbackText: optionalString(value.fallbackText)
    };
  }

  if (value.kind === "map_action") {
    if (typeof value.action !== "string" || value.surface !== "map_overlay" || !toTrustedMapAction(value.action, value.params)) {
      return null;
    }

    return {
      kind: "map_action",
      schemaVersion: "agent-ui@1",
      action: value.action,
      surface: "map_overlay",
      params: value.params,
      fallbackText: optionalString(value.fallbackText)
    };
  }

  return null;
}

export function isUiEnvelopeAllowed(envelope: UIEnvelope, registry: UIRegistry | null) {
  if (!registry) {
    return false;
  }

  if (envelope.kind === "registered_component") {
    if (!LOCAL_COMPONENT_SURFACES.get(envelope.component)?.includes(envelope.surface)) return false;
    const manifest = registry.components.find((item) => item.id === envelope.component);
    return Boolean(manifest?.supportedSurfaces.includes(envelope.surface));
  }

  if (envelope.kind === "chart") {
    return envelope.grammar === "echarts-safe-subset" && registry.chart_grammars.includes(envelope.grammar);
  }

  if (!LOCAL_ACTION_SURFACES.get(envelope.action)?.includes(envelope.surface)) return false;
  const manifest = registry.actions.find((item) => item.id === envelope.action);
  return Boolean(manifest?.supportedSurfaces.includes(envelope.surface));
}

export function parseUiRegistry(value: unknown): UIRegistry | null {
  if (!isRecord(value) || value.schema_version !== "agent-ui-registry@1") {
    return null;
  }

  return {
    schema_version: "agent-ui-registry@1",
    chart_grammars: stringArray(value.chart_grammars),
    components: manifestArray(value.components),
    actions: manifestArray(value.actions)
  };
}

function manifestArray(value: unknown) {
  if (!Array.isArray(value)) {
    return [];
  }

  return value.flatMap((item) => {
    if (!isRecord(item) || typeof item.id !== "string") {
      return [];
    }

    return [{ id: item.id, supportedSurfaces: stringArray(item.supportedSurfaces).filter(isSurface) }];
  });
}

function stringArray(value: unknown) {
  return Array.isArray(value) ? value.filter((item): item is string => typeof item === "string") : [];
}

function isChartSurface(value: unknown): value is "chat_inline" | "side_panel" | "canvas" {
  return value === "chat_inline" || value === "side_panel" || value === "canvas";
}

function isSurface(value: unknown): value is UISurface {
  return typeof value === "string" && surfaces.has(value as UISurface);
}

function optionalString(value: unknown) {
  return typeof value === "string" ? value : undefined;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

function boundedString(value: unknown, maxLength: number) {
  return typeof value === "string" && value.trim().length > 0 && value.length <= maxLength ? value.trim() : null;
}

function isFiniteTimestamp(value: unknown): value is number {
  return typeof value === "number" && Number.isFinite(value) && value > 0;
}

function hasOnlyKeys(value: Record<string, unknown>, allowed: readonly string[]) {
  const keys = new Set(allowed);
  return Object.keys(value).every((key) => keys.has(key));
}

function isValidComponentProps(component: string, props: unknown) {
  if (!isRecord(props)) return false;
  if (component === "HistoryPanel") {
    if (!hasOnlyKeys(props, ["reason", "feature_infos", "data_type", "start_time", "end_time"])) return false;
    const items = props.feature_infos;
    return Array.isArray(items) && items.length > 0 && items.length <= 100 &&
      items.every((item) => Array.isArray(item) && item.length === 2 && item.every((part) => boundedString(part, 128))) &&
      (props.data_type === "realtime" || props.data_type === "scheme" || props.data_type === "none");
  }
  return false;
}

function isValidChart(spec: unknown, data: unknown) {
  if (!isRecord(spec) || !hasOnlyKeys(spec, ["title", "chart_type", "x_axis_name", "y_axis_name"])) return false;
  if (spec.chart_type !== undefined && spec.chart_type !== "line" && spec.chart_type !== "bar") return false;
  const parsedSpec = parseChartSpec(spec);
  const parsedData = normalizeSafeChartData(data, 100);
  return (parsedSpec.chartType === "line" || parsedSpec.chartType === "bar") && parsedData !== null &&
    parsedData.xData.length <= 100 && parsedData.series.length <= 8 &&
    parsedData.series.every((series) => series.data.length === parsedData.xData.length);
}