agent-r/src/output-schema.ts

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

function hasObjectType(schema: Record<string, unknown>): boolean {
  const { type } = schema;
  if (type === "object") {
    return true;
  }
  return Array.isArray(type) && type.includes("object");
}

function hasArrayType(schema: Record<string, unknown>): boolean {
  const { type } = schema;
  if (type === "array") {
    return true;
  }
  return Array.isArray(type) && type.includes("array");
}

type NormalizeResult =
  | {
      ok: true;
      schema: unknown;
    }
  | {
      ok: false;
      reason: string;
    };

const UNSUPPORTED_OBJECT_KEYWORDS = [
  "patternProperties",
  "propertyNames",
  "dependentSchemas",
  "unevaluatedProperties",
  "allOf",
  "anyOf",
  "oneOf",
  "if",
  "then",
  "else",
  "not",
];

function normalizeSchemaNode(node: unknown, path: string): NormalizeResult {
  if (!isPlainObject(node)) {
    return { ok: true, schema: node };
  }

  for (const keyword of UNSUPPORTED_OBJECT_KEYWORDS) {
    if (keyword in node) {
      return {
        ok: false,
        reason: `Structured outputs do not safely support '${keyword}' at ${path}.`,
      };
    }
  }

  const normalized: Record<string, unknown> = { ...node };

  if (hasObjectType(node) || "properties" in node || "additionalProperties" in node) {
    const propertiesValue = node.properties;
    if (propertiesValue !== undefined && !isPlainObject(propertiesValue)) {
      return {
        ok: false,
        reason: `Object schema at ${path} must declare properties as an object.`,
      };
    }

    const properties = isPlainObject(propertiesValue) ? propertiesValue : {};
    const normalizedProperties: Record<string, unknown> = {};
    for (const [key, value] of Object.entries(properties)) {
      const child = normalizeSchemaNode(value, `${path}.properties.${key}`);
      if (!child.ok) {
        return child;
      }
      normalizedProperties[key] = child.schema;
    }

    const additionalProperties = node.additionalProperties;
    if (additionalProperties !== false) {
      return {
        ok: false,
        reason: `Structured outputs require additionalProperties: false at ${path}.`,
      };
    }

    const propertyKeys = Object.keys(normalizedProperties);
    const required = node.required;
    if (
      !Array.isArray(required) ||
      required.length !== propertyKeys.length ||
      required.some((value) => typeof value !== "string") ||
      propertyKeys.some((key) => !required.includes(key)) ||
      required.some((key) => !propertyKeys.includes(key))
    ) {
      return {
        ok: false,
        reason: `Structured outputs require 'required' to match every property at ${path}.`,
      };
    }

    normalized.properties = normalizedProperties;
    normalized.required = required;
    normalized.additionalProperties = additionalProperties;
  }

  if (hasArrayType(node) || "items" in node) {
    if (!("items" in node)) {
      return {
        ok: false,
        reason: `Array schema at ${path} is missing 'items'.`,
      };
    }
    if (Array.isArray(node.items)) {
      return {
        ok: false,
        reason: `Structured outputs do not safely support tuple-style 'items' at ${path}.`,
      };
    }
    const normalizedItems = normalizeSchemaNode(node.items, `${path}.items`);
    if (!normalizedItems.ok) {
      return normalizedItems;
    }
    normalized.items = normalizedItems.schema;
  }

  return {
    ok: true,
    schema: normalized,
  };
}

export function prepareOutputSchema(schema: Record<string, unknown>): {
  outputSchema: Record<string, unknown> | null;
  fallbackReason: string | null;
} {
  const normalized = normalizeSchemaNode(schema, "$");
  if (!normalized.ok) {
    return {
      outputSchema: null,
      fallbackReason: normalized.reason,
    };
  }

  if (!isPlainObject(normalized.schema)) {
    return {
      outputSchema: null,
      fallbackReason: "Structured outputs require a plain JSON object schema.",
    };
  }

  return {
    outputSchema: normalized.schema,
    fallbackReason: null,
  };
}

export function buildSchemaFallbackPrompt(
  prompt: string,
  schema: Record<string, unknown>,
  reason?: string,
): string {
  return [
    prompt,
    "Structured output enforcement is unavailable for this turn.",
    reason ? `Reason: ${reason}` : null,
    "Return a single JSON object only. Do not include markdown fences or any commentary.",
    "Match this JSON schema exactly:",
    JSON.stringify(schema, null, 2),
  ]
    .filter((value): value is string => Boolean(value))
    .join("\n\n");
}

export function buildValidationRetryPrompt(
  message: string,
  errorText: string,
  schema: Record<string, unknown>,
  schemaEmbeddedInPrompt: boolean,
): string {
  return [
    "The previous response did not parse or validate.",
    `Problem: ${errorText}`,
    schemaEmbeddedInPrompt
      ? [
          "Respond again with JSON only.",
          "Match this JSON schema exactly:",
          JSON.stringify(schema, null, 2),
        ].join("\n\n")
      : "Respond again with JSON only, matching the schema exactly.",
    "Previous response:",
    message,
  ].join("\n\n");
}

export function isInvalidOutputSchemaError(message: string): boolean {
  const lowered = message.toLowerCase();
  return (
    lowered.includes("invalid_json_schema") ||
    lowered.includes("invalid schema for response_format") ||
    lowered.includes("text.format.schema")
  );
}