ThothII/backend/src/workspaces/schema.ts

import { parseAllDocuments, stringify } from "yaml";
import { z } from "zod";

export const DWH_TRANSPORTS = ["postgres_direct", "rest_api", "ssh_tunnel"] as const;
export type DwhTransport = (typeof DWH_TRANSPORTS)[number];

export const VECTOR_TRANSPORTS = ["pgvector_direct", "rest_api", "ssh_tunnel"] as const;
export type VectorTransport = (typeof VECTOR_TRANSPORTS)[number];

export const REST_DIAGNOSTIC_METHODS = ["GET", "POST"] as const;
export type RestDiagnosticMethod = (typeof REST_DIAGNOSTIC_METHODS)[number];

export const DIAGNOSTIC_AUTH_MODES = ["none", "bearer"] as const;
export type DiagnosticAuthMode = (typeof DIAGNOSTIC_AUTH_MODES)[number];

export interface RestDiagnosticRequest {
  method: RestDiagnosticMethod;
  path: string;
  auth: DiagnosticAuthMode;
}

export interface CanonicalDiagnostics {
  dwh_rest?: RestDiagnosticRequest & { response: { database: string; schema: string } };
  vector_rest?: {
    metadata: RestDiagnosticRequest & {
      response: { collection: string; dimensions: string; distance: string };
    };
    reversible_probe?: RestDiagnosticRequest & { method: "POST" };
  };
  embedding?: RestDiagnosticRequest & { response: { model: string; dimensions: string } };
}

interface WorkspaceMetadata<Version extends 1 | 2> {
  schema_version: Version;
  id: string;
  name: string;
  description?: string;
  language: "en" | "it";
}

interface WorkspaceDwh {
  engine: "postgres";
  database: string;
  schema: string;
  port?: number;
  timeout_ms?: number;
  supported_transports: DwhTransport[];
}

interface VectorStore {
  engine: "pgvector";
  collection: string;
  dimensions: number;
  distance: "cosine" | "l2" | "inner_product";
  port?: number;
  timeout_ms?: number;
  supported_transports: VectorTransport[];
}

interface SemanticIndex<TVectorStore extends VectorStore> {
  vector_store: TVectorStore;
  vector_writer?: Record<string, never>;
  embedding: {
    provider: "ollama_compatible" | "openai_compatible";
    model: string;
    dimensions: number;
    timeout_ms?: number;
  };
}

interface WorkspaceBase<Version extends 1 | 2, TVectorStore extends VectorStore> {
  workspace: WorkspaceMetadata<Version>;
  dwh: WorkspaceDwh;
  semantic_index: SemanticIndex<TVectorStore>;
  llm_policy: {
    default?: `${string}/${string}`;
    allowed: `${string}/${string}`[];
  };
  diagnostics?: CanonicalDiagnostics;
}

export interface CanonicalWorkspace extends WorkspaceBase<2, VectorStore & { database: string; schema: string }> {}

/** A readable, non-operational v1 descriptor. It must be explicitly migrated before use. */
export interface LegacyWorkspace extends WorkspaceBase<1, VectorStore & { database?: string; schema?: string }> {}

export type WorkspaceDescriptor = CanonicalWorkspace | LegacyWorkspace;

const workspaceId = z.string().regex(/^[a-z][a-z0-9-]{2,62}$/, {
  message: "workspace id must match ^[a-z][a-z0-9-]{2,62}$",
});
const identifier = z.string().regex(/^[A-Za-z_][A-Za-z0-9_]*$/, {
  message: "database identifiers must start with a letter or underscore",
});
const dimensions = z.number().int().positive().max(32_768);
const port = z.number().int().min(1).max(65_535);
const timeoutMs = z.number().int().positive();
const modelReference = z.string().regex(/^[^/\s]+\/[^/\s]+$/, {
  message: "model must use provider/model syntax",
});

function isOriginRelativeDiagnosticPath(value: string): boolean {
  return /^\/(?!\/)[^\\\u0000-\u001F\u007F?#]*$/.test(value) && !/%5c/i.test(value);
}

const diagnosticPath = z.string().refine(isOriginRelativeDiagnosticPath, {
  message: "diagnostic paths must be origin-relative and cannot contain backslashes, control characters, queries, or fragments",
});
const responseField = z.string().regex(/^[A-Za-z_][A-Za-z0-9_]*$/, {
  message: "diagnostic response fields must be identifiers",
});
const restDiagnosticRequest = z.object({
  method: z.enum(REST_DIAGNOSTIC_METHODS),
  path: diagnosticPath,
  auth: z.enum(DIAGNOSTIC_AUTH_MODES),
}).strict();
const dwhRestDiagnostic = restDiagnosticRequest.extend({
  response: z.object({ database: responseField, schema: responseField }).strict(),
}).strict();
const vectorMetadataDiagnostic = restDiagnosticRequest.extend({
  response: z.object({
    collection: responseField,
    dimensions: responseField,
    distance: responseField,
  }).strict(),
}).strict();
const reversibleVectorProbe = restDiagnosticRequest.extend({ method: z.literal("POST") }).strict();
const embeddingDiagnostic = restDiagnosticRequest.extend({
  response: z.object({ model: responseField, dimensions: responseField }).strict(),
}).strict();
const diagnosticsSchema = z.object({
  dwh_rest: dwhRestDiagnostic.optional(),
  vector_rest: z.object({
    metadata: vectorMetadataDiagnostic,
    reversible_probe: reversibleVectorProbe.optional(),
  }).strict().optional(),
  embedding: embeddingDiagnostic.optional(),
}).strict().optional();

const dwhSchema = z.object({
  engine: z.literal("postgres"),
  database: identifier,
  schema: identifier,
  port: port.optional(),
  timeout_ms: timeoutMs.optional(),
  supported_transports: z.array(z.enum(DWH_TRANSPORTS)).min(1),
}).strict();
const embeddingSchema = z.object({
  provider: z.enum(["ollama_compatible", "openai_compatible"]),
  model: z.string().trim().min(1),
  dimensions,
  timeout_ms: timeoutMs.optional(),
}).strict();
const vectorStoreShape = {
  engine: z.literal("pgvector"),
  collection: identifier,
  dimensions,
  distance: z.enum(["cosine", "l2", "inner_product"]),
  port: port.optional(),
  timeout_ms: timeoutMs.optional(),
  supported_transports: z.array(z.enum(VECTOR_TRANSPORTS)).min(1),
};
const legacyVectorStoreSchema = z.object({
  ...vectorStoreShape,
  database: identifier.optional(),
  schema: identifier.optional(),
}).strict();
const canonicalVectorStoreSchema = z.object({
  ...vectorStoreShape,
  database: identifier,
  schema: identifier,
}).strict();
const llmPolicySchema = z.object({
  default: modelReference.optional(),
  allowed: z.array(modelReference).min(1),
}).strict();

function unique<T>(values: readonly T[], context: z.RefinementCtx, path: PropertyKey[]) {
  if (new Set(values).size !== values.length) {
    context.addIssue({ code: "custom", path, message: "supported transports must not repeat" });
  }
}

function workspaceInvariants(workspace: any, context: z.RefinementCtx): void {
  unique(workspace.dwh.supported_transports, context, ["dwh", "supported_transports"]);
  unique(
    workspace.semantic_index.vector_store.supported_transports,
    context,
    ["semantic_index", "vector_store", "supported_transports"],
  );
  unique(workspace.llm_policy.allowed, context, ["llm_policy", "allowed"]);

  if (workspace.semantic_index.vector_store.dimensions !== workspace.semantic_index.embedding.dimensions) {
    context.addIssue({
      code: "custom",
      path: ["semantic_index", "embedding", "dimensions"],
      message: "embedding dimensions must match vector store dimensions",
    });
  }
  if (workspace.llm_policy.default && !workspace.llm_policy.allowed.includes(workspace.llm_policy.default)) {
    context.addIssue({
      code: "custom",
      path: ["llm_policy", "default"],
      message: "LLM default must be included in the allowlist",
    });
  }
  if (workspace.diagnostics?.dwh_rest && !workspace.dwh.supported_transports.includes("rest_api")) {
    context.addIssue({
      code: "custom",
      path: ["diagnostics", "dwh_rest"],
      message: "diagnostics.dwh_rest requires dwh rest_api transport support",
    });
  }
  if (
    workspace.diagnostics?.vector_rest
    && !workspace.semantic_index.vector_store.supported_transports.includes("rest_api")
  ) {
    context.addIssue({
      code: "custom",
      path: ["diagnostics", "vector_rest"],
      message: "diagnostics.vector_rest requires vector_store rest_api transport support",
    });
  }
}

const workspaceShape = {
  dwh: dwhSchema,
  llm_policy: llmPolicySchema,
  diagnostics: diagnosticsSchema,
};
const LegacyWorkspaceSchema = z.object({
  ...workspaceShape,
  workspace: z.object({
    schema_version: z.literal(1), id: workspaceId, name: z.string().trim().min(1),
    description: z.string().trim().min(1).optional(), language: z.enum(["en", "it"]),
  }).strict(),
  semantic_index: z.object({
    vector_store: legacyVectorStoreSchema,
    vector_writer: z.object({}).strict().optional(),
    embedding: embeddingSchema,
  }).strict(),
}).strict().superRefine(workspaceInvariants);
const CanonicalWorkspaceSchema = z.object({
  ...workspaceShape,
  workspace: z.object({
    schema_version: z.literal(2), id: workspaceId, name: z.string().trim().min(1),
    description: z.string().trim().min(1).optional(), language: z.enum(["en", "it"]),
  }).strict(),
  semantic_index: z.object({
    vector_store: canonicalVectorStoreSchema,
    vector_writer: z.object({}).strict().optional(),
    embedding: embeddingSchema,
  }).strict(),
}).strict().superRefine(workspaceInvariants);
const WorkspaceDescriptorSchema = z.union([CanonicalWorkspaceSchema, LegacyWorkspaceSchema]);

export function parseWorkspaceYaml(source: string): WorkspaceDescriptor {
  const documents = parseAllDocuments(source, { uniqueKeys: true });
  if (documents.length !== 1) throw new Error("Workspace YAML must contain exactly one document");
  const document = documents[0];
  if (document.errors.length > 0) {
    throw new Error(`Invalid workspace YAML: ${document.errors.map((error) => error.message).join("; ")}`);
  }
  return validateWorkspaceDescriptor(document.toJSON());
}

export function validateWorkspaceDescriptor(workspace: unknown): WorkspaceDescriptor {
  return WorkspaceDescriptorSchema.parse(workspace) as WorkspaceDescriptor;
}

export function isCanonicalWorkspace(workspace: WorkspaceDescriptor): workspace is CanonicalWorkspace {
  return workspace.workspace.schema_version === 2;
}

/** Rejects readable v1 descriptors at every operational boundary until a caller migrates them. */
export function validateCanonicalWorkspace(workspace: unknown): CanonicalWorkspace {
  const descriptor = validateWorkspaceDescriptor(workspace);
  if (!isCanonicalWorkspace(descriptor)) {
    throw new Error("Workspace descriptor requires explicit migration to schema version 2");
  }
  return descriptor;
}

/**
 * Explicitly upgrades a readable v1 descriptor. The caller must supply vector identity; the
 * transformer never derives it from DWH identity, even where both services share a database.
 */
export function migrateWorkspaceV1ToV2(
  workspace: LegacyWorkspace,
  vectorIdentity: { database: string; schema: string },
): CanonicalWorkspace {
  const legacy = LegacyWorkspaceSchema.parse(workspace) as LegacyWorkspace;
  const identity = z.object({ database: identifier, schema: identifier }).strict().parse(vectorIdentity);
  return validateCanonicalWorkspace({
    ...legacy,
    workspace: { ...legacy.workspace, schema_version: 2 },
    semantic_index: {
      ...legacy.semantic_index,
      vector_store: { ...legacy.semantic_index.vector_store, ...identity },
    },
  });
}

/** Builds a request URL only after rejecting values that can leave the declared service origin. */
export function resolveDiagnosticUrl(baseUrl: string, path: string): URL {
  if (!isOriginRelativeDiagnosticPath(path)) throw new Error("Diagnostic path must remain on the configured origin");
  const base = new URL(baseUrl);
  const resolved = new URL(path, base);
  if (resolved.origin !== base.origin) throw new Error("Diagnostic URL must remain on the configured origin");
  return resolved;
}

export function serializeWorkspaceYaml(workspace: CanonicalWorkspace): string {
  const canonical = validateCanonicalWorkspace(workspace);
  return stringify(canonical, { lineWidth: 0, sortMapEntries: true });
}

export { buildInstallationContract, renderWorkspaceDocs } from "./contracts.js";