DiAL/src/server/checker/normalizer.ts

import { isPlainObject } from "es-toolkit";

import type { ConfigValidationIssue } from "./schema/issues";
import type { AuthoringProbeConfig, NormalizedProbeConfig } from "./schema/types";
import type { RawTargetConfig } from "./types";

import { resolveContentExpectations } from "./expect/content";
import { CONTENT_EXTRACTOR_KEY_SET, MATCHER_KEY_SET } from "./expect/keys";
import { isValueMatcherObject, isValueMatcherPrimitive, resolveValueExpectation } from "./expect/value";
import { resolveVariables } from "./variables";

type ExpectRecord = Record<string, unknown>;

export function normalizeAuthoringConfig(config: unknown): {
  config: unknown;
  issues: ConfigValidationIssue[];
} {
  const variableResult = resolveVariables(config);
  if (!isPlainObject(variableResult.config)) {
    return variableResult;
  }

  const normalized = { ...(variableResult.config as Record<string, unknown>) };
  delete normalized["variables"];
  if (Array.isArray(normalized["targets"])) {
    normalized["targets"] = normalized["targets"].map((target) => normalizeTarget(target));
  }

  return { config: normalized, issues: variableResult.issues };
}

function canNormalizeContentEntry(value: unknown): boolean {
  if (!isPlainObject(value)) return false;
  const keys = Object.keys(value);
  const extractorKeys = keys.filter((key) => CONTENT_EXTRACTOR_KEY_SET.has(key));
  const matcherKeys = keys.filter((key) => MATCHER_KEY_SET.has(key));
  if (extractorKeys.length === 0) return matcherKeys.length > 0 && matcherKeys.length === keys.length;
  if (extractorKeys.length !== 1 || matcherKeys.length > 0 || keys.length !== 1) return false;
  return isPlainObject((value as ExpectRecord)[extractorKeys[0]!]);
}

function compact(original: ExpectRecord, overrides: ExpectRecord): ExpectRecord {
  const result: ExpectRecord = {};
  for (const [key, value] of Object.entries(original)) {
    if (value !== undefined) result[key] = value;
  }
  for (const [key, value] of Object.entries(overrides)) {
    if (value !== undefined) result[key] = value;
  }
  return result;
}

function normalizeCommandExpect(raw: ExpectRecord): ExpectRecord {
  return compact(raw, {
    durationMs: normalizeValue(raw["durationMs"]),
    exitCode: raw["exitCode"],
    stderr: normalizeContent(raw["stderr"]),
    stdout: normalizeContent(raw["stdout"]),
  });
}

function normalizeContent(value: unknown): unknown {
  if (value === undefined) return undefined;
  if (!Array.isArray(value)) return value;
  return (value as unknown[]).map((entry): unknown => {
    if (!canNormalizeContentEntry(entry)) return entry;
    const resolved = resolveContentExpectations([entry] as never);
    return resolved?.[0];
  });
}

function normalizeDbExpect(raw: ExpectRecord): ExpectRecord {
  return compact(raw, {
    durationMs: normalizeValue(raw["durationMs"]),
    result: normalizeContent(raw["result"]),
    rowCount: normalizeValue(raw["rowCount"]),
    rows: Array.isArray(raw["rows"]) ? raw["rows"].map((row) => normalizeKeyed(row)) : raw["rows"],
  });
}

function normalizeExpect(type: string, expect: unknown): unknown {
  if (!isPlainObject(expect)) return expect;
  const raw = expect as ExpectRecord;
  switch (type) {
    case "cmd":
      return normalizeCommandExpect(raw);
    case "db":
      return normalizeDbExpect(raw);
    case "http":
      return normalizeHttpExpect(raw);
    case "icmp":
      return normalizeIcmpExpect(raw);
    case "llm":
      return normalizeLlmExpect(raw);
    case "tcp":
      return normalizeTcpExpect(raw);
    case "udp":
      return normalizeUdpExpect(raw);
    default:
      return expect;
  }
}

function normalizeHttpExpect(raw: ExpectRecord): ExpectRecord {
  return compact(raw, {
    body: normalizeContent(raw["body"]),
    durationMs: normalizeValue(raw["durationMs"]),
    headers: normalizeKeyed(raw["headers"]),
    status: raw["status"],
  });
}

function normalizeIcmpExpect(raw: ExpectRecord): ExpectRecord {
  return compact(raw, {
    alive: raw["alive"],
    avgLatencyMs: normalizeValue(raw["avgLatencyMs"]),
    durationMs: normalizeValue(raw["durationMs"]),
    maxLatencyMs: normalizeValue(raw["maxLatencyMs"]),
    packetLossPercent: normalizeValue(raw["packetLossPercent"]),
  });
}

function normalizeKeyed(value: unknown): unknown {
  if (value === undefined) return undefined;
  if (!isPlainObject(value)) return value;
  return Object.entries(value as ExpectRecord).map(([key, item]) => ({ key, matcher: normalizeValue(item) }));
}

function normalizeLlmExpect(raw: ExpectRecord): ExpectRecord {
  return compact(raw, {
    durationMs: normalizeValue(raw["durationMs"]),
    finishReason: normalizeValue(raw["finishReason"]),
    headers: normalizeKeyed(raw["headers"]),
    output: normalizeContent(raw["output"]),
    rawFinishReason: normalizeValue(raw["rawFinishReason"]),
    status: raw["status"],
    stream: isPlainObject(raw["stream"])
      ? compact(raw["stream"] as ExpectRecord, {
          completed: (raw["stream"] as ExpectRecord)["completed"],
          firstTokenMs: normalizeValue((raw["stream"] as ExpectRecord)["firstTokenMs"]),
        })
      : raw["stream"],
    usage: isPlainObject(raw["usage"])
      ? compact(raw["usage"] as ExpectRecord, {
          inputTokens: normalizeValue((raw["usage"] as ExpectRecord)["inputTokens"]),
          outputTokens: normalizeValue((raw["usage"] as ExpectRecord)["outputTokens"]),
          totalTokens: normalizeValue((raw["usage"] as ExpectRecord)["totalTokens"]),
        })
      : raw["usage"],
  });
}

function normalizeTarget(target: unknown): unknown {
  if (!isPlainObject(target)) return target;
  const result = { ...(target as RawTargetConfig) };
  if (result.expect !== undefined) {
    result.expect = normalizeExpect(result.type, result.expect);
  }
  return result;
}

function normalizeTcpExpect(raw: ExpectRecord): ExpectRecord {
  return compact(raw, {
    banner: normalizeContent(raw["banner"]),
    connected: raw["connected"],
    durationMs: normalizeValue(raw["durationMs"]),
  });
}

function normalizeUdpExpect(raw: ExpectRecord): ExpectRecord {
  return compact(raw, {
    durationMs: normalizeValue(raw["durationMs"]),
    responded: raw["responded"],
    response: normalizeContent(raw["response"]),
    responseSize: normalizeValue(raw["responseSize"]),
    sourceHost: normalizeValue(raw["sourceHost"]),
    sourcePort: normalizeValue(raw["sourcePort"]),
  });
}

function normalizeValue(value: unknown): unknown {
  if (value === undefined) return undefined;
  if (isValueMatcherPrimitive(value) || isValueMatcherObject(value)) return resolveValueExpectation(value);
  return value;
}

export type { AuthoringProbeConfig, NormalizedProbeConfig };