ThothII/backend/test/workspaces-diagnostics.test.ts

import { expect, test, vi } from "vitest";
import {
  createProductionWorkspaceDiagnoser,
  createWorkspaceDiagnoser,
  type DiagnosticAdapters,
} from "../src/workspaces/diagnostics.js";
import type { RuntimeBindings } from "../src/workspaces/runtime-renderer.js";
import { parseWorkspaceYaml } from "../src/workspaces/schema.js";

const workspace = parseWorkspaceYaml(`workspace:
  schema_version: 1
  id: psd-clinical
  name: Policlinico San Donato
  language: it
dwh:
  engine: postgres
  database: warehouse
  schema: datawarehouse
  timeout_ms: 8000
  supported_transports: [postgres_direct, rest_api, ssh_tunnel]
semantic_index:
  vector_store:
    engine: pgvector
    collection: clinical_documents
    dimensions: 768
    distance: cosine
    timeout_ms: 8000
    supported_transports: [pgvector_direct, rest_api, ssh_tunnel]
  embedding:
    provider: ollama_compatible
    model: nomic-embed-text-v2-moe
    dimensions: 768
    timeout_ms: 8000
llm_policy:
  allowed: [zai/glm-5.2]
`);

const bindings: RuntimeBindings = {
  dwh: {
    transport: "postgres_direct",
    missing: [],
    values: {
      THT_WS_PSD_CLINICAL_DWH_HOST: "dwh.example.test",
      THT_WS_PSD_CLINICAL_DWH_PORT: "5432",
      THT_WS_PSD_CLINICAL_DWH_USER: "reader",
      THT_WS_PSD_CLINICAL_DWH_PASSWORD_FILE: "/run/secrets/dwh-password",
      THT_WS_PSD_CLINICAL_DWH_TLS_CA_FILE: "/run/secrets/dwh-ca",
    },
  },
  vector: {
    transport: "pgvector_direct",
    missing: [],
    values: {
      THT_WS_PSD_CLINICAL_VECTOR_HOST: "vector.example.test",
      THT_WS_PSD_CLINICAL_VECTOR_PORT: "5432",
      THT_WS_PSD_CLINICAL_VECTOR_USER: "vector-reader",
      THT_WS_PSD_CLINICAL_VECTOR_PASSWORD_FILE: "/run/secrets/vector-password",
      THT_WS_PSD_CLINICAL_VECTOR_TLS_CA_FILE: "/run/secrets/vector-ca",
    },
  },
  embedding: {
    transport: "rest_api",
    missing: [],
    values: {
      THT_WS_PSD_CLINICAL_EMBEDDING_BASE_URL: "https://embedding.example.test",
      THT_WS_PSD_CLINICAL_EMBEDDING_API_KEY_FILE: "/run/secrets/embedding-key",
      THT_WS_PSD_CLINICAL_EMBEDDING_TLS_CA_FILE: "/run/secrets/embedding-ca",
    },
  },
};

function successfulAdapters(overrides: Partial<DiagnosticAdapters> = {}): DiagnosticAdapters {
  return {
    probeConnector: vi.fn(async (request) => ({
      resolved: true,
      tlsVerified: true,
      authenticated: true,
      resource: request.resource,
    })),
    withSshTunnel: vi.fn(async (_request, probe) => probe({ host: "127.0.0.1", port: 45678 })),
    inspectVector: vi.fn(async () => ({
      collection: "clinical_documents",
      dimensions: 768,
      distance: "cosine",
    })),
    probeEmbedding: vi.fn(async () => ({ available: true, dimensions: 768 })),
    writeDiagnosticRecord: vi.fn(async () => undefined),
    removeDiagnosticRecord: vi.fn(async () => undefined),
    ...overrides,
  };
}

function diagnose(adapters = successfulAdapters()) {
  return createWorkspaceDiagnoser(adapters, { timeoutMs: 5000 });
}

test("reports the missing vector collection dimensions as semantic-index incompatibility", async () => {
  const result = await diagnose(successfulAdapters({
    inspectVector: vi.fn(async () => ({
      collection: "clinical_documents",
      dimensions: undefined,
      distance: "cosine",
    })),
  }))(workspace, bindings, { writeProbe: false });

  expect(result.diagnostics).toContainEqual(expect.objectContaining({
    code: "semantic_index_incompatible",
  }));
  expect(result.activatable).toBe(false);
});

test("refuses an SSH tunnel when known-hosts is missing", async () => {
  const sshBindingsWithoutKnownHosts: RuntimeBindings = {
    ...bindings,
    dwh: {
      transport: "ssh_tunnel",
      missing: ["THT_WS_PSD_CLINICAL_DWH_SSH_KNOWN_HOSTS_FILE"],
      values: {
        THT_WS_PSD_CLINICAL_DWH_USER: "reader",
        THT_WS_PSD_CLINICAL_DWH_PASSWORD_FILE: "/run/secrets/dwh-password",
        THT_WS_PSD_CLINICAL_DWH_SSH_HOST: "bastion.example.test",
        THT_WS_PSD_CLINICAL_DWH_SSH_PORT: "22",
        THT_WS_PSD_CLINICAL_DWH_SSH_USER: "tunnel",
        THT_WS_PSD_CLINICAL_DWH_SSH_PRIVATE_KEY_FILE: "/run/secrets/ssh-key",
        THT_WS_PSD_CLINICAL_DWH_SSH_TARGET_HOST: "dwh.internal",
        THT_WS_PSD_CLINICAL_DWH_SSH_TARGET_PORT: "5432",
      },
    },
  };
  const adapters = successfulAdapters();

  const result = await diagnose(adapters)(workspace, sshBindingsWithoutKnownHosts, { writeProbe: false });

  expect(result.activatable).toBe(false);
  expect(result.diagnostics[0]).toMatchObject({
    code: "binding_missing",
    field: expect.stringContaining("SSH_KNOWN_HOSTS_FILE"),
  });
  expect(adapters.withSshTunnel).not.toHaveBeenCalled();
});

test("checks direct and REST resolution, TLS, authentication, and resource metadata without exposing failures", async () => {
  const adapters = successfulAdapters({
    probeConnector: vi.fn(async (request) => ({
      resolved: true,
      tlsVerified: true,
      authenticated: true,
      resource: request.role === "dwh"
        ? { database: "warehouse", schema: "wrong_schema" }
        : request.resource,
    })),
  });
  const restBindings: RuntimeBindings = {
    ...bindings,
    dwh: {
      transport: "rest_api",
      missing: [],
      values: {
        THT_WS_PSD_CLINICAL_DWH_BASE_URL: "https://dwh.example.test",
        THT_WS_PSD_CLINICAL_DWH_API_KEY_FILE: "/run/secrets/dwh-api-key",
        THT_WS_PSD_CLINICAL_DWH_TLS_CA_FILE: "/run/secrets/dwh-ca",
      },
    },
  };

  const result = await diagnose(adapters)(workspace, restBindings, { writeProbe: false });

  expect(adapters.probeConnector).toHaveBeenCalledWith(expect.objectContaining({
    transport: "rest_api",
    timeoutMs: 5000,
    tlsCaFile: "/run/secrets/dwh-ca",
    credentialFile: "/run/secrets/dwh-api-key",
    resource: { database: "warehouse", schema: "datawarehouse" },
  }));
  expect(result).toMatchObject({ activatable: false });
  expect(JSON.stringify(result)).not.toContain("wrong_schema");
  expect(JSON.stringify(result)).not.toContain("/run/secrets/dwh-api-key");
});

test("uses a loopback-only SSH tunnel for the bounded connector probe", async () => {
  const adapters = successfulAdapters();
  const sshBindings: RuntimeBindings = {
    ...bindings,
    dwh: {
      transport: "ssh_tunnel",
      missing: [],
      values: {
        THT_WS_PSD_CLINICAL_DWH_USER: "reader",
        THT_WS_PSD_CLINICAL_DWH_PASSWORD_FILE: "/run/secrets/dwh-password",
        THT_WS_PSD_CLINICAL_DWH_TLS_CA_FILE: "/run/secrets/dwh-ca",
        THT_WS_PSD_CLINICAL_DWH_SSH_HOST: "bastion.example.test",
        THT_WS_PSD_CLINICAL_DWH_SSH_PORT: "22",
        THT_WS_PSD_CLINICAL_DWH_SSH_USER: "tunnel",
        THT_WS_PSD_CLINICAL_DWH_SSH_PRIVATE_KEY_FILE: "/run/secrets/ssh-key",
        THT_WS_PSD_CLINICAL_DWH_SSH_KNOWN_HOSTS_FILE: "/run/secrets/known-hosts",
        THT_WS_PSD_CLINICAL_DWH_SSH_TARGET_HOST: "dwh.internal",
        THT_WS_PSD_CLINICAL_DWH_SSH_TARGET_PORT: "5432",
      },
    },
  };

  const result = await diagnose(adapters)(workspace, sshBindings, { writeProbe: false });

  expect(result.activatable).toBe(true);
  expect(adapters.withSshTunnel).toHaveBeenCalledWith(expect.objectContaining({
    localHost: "127.0.0.1",
    localPort: 0,
    knownHostsFile: "/run/secrets/known-hosts",
    timeoutMs: 5000,
  }), expect.any(Function));
  expect(adapters.probeConnector).toHaveBeenCalledWith(expect.objectContaining({
    host: "127.0.0.1",
    port: 45678,
  }));
});

test("requires a matching embedding model vector and removes its unique write probe", async () => {
  const adapters = successfulAdapters();

  const result = await diagnose(adapters)(workspace, bindings, { writeProbe: true });

  expect(result.activatable).toBe(true);
  expect(adapters.probeEmbedding).toHaveBeenCalledWith(expect.objectContaining({
    model: "nomic-embed-text-v2-moe",
    timeoutMs: 5000,
  }));
  expect(adapters.writeDiagnosticRecord).toHaveBeenCalledWith(expect.objectContaining({
    collection: "clinical_documents",
    id: expect.stringMatching(/^diagnostic:/),
    dimensions: 768,
  }));
  expect(adapters.removeDiagnosticRecord).toHaveBeenCalledWith(expect.objectContaining({
    collection: "clinical_documents",
    id: expect.stringMatching(/^diagnostic:/),
  }));
});

test("constructs the production diagnoser with the configured timeout and injected adapters", async () => {
  const adapters = successfulAdapters();

  const result = await createProductionWorkspaceDiagnoser(1234, adapters)(workspace, bindings, {
    writeProbe: false,
  });

  expect(result.activatable).toBe(true);
  expect(adapters.probeConnector).toHaveBeenCalledWith(expect.objectContaining({ timeoutMs: 1234 }));
  expect(adapters.probeEmbedding).toHaveBeenCalledWith(expect.objectContaining({ timeoutMs: 1234 }));
});

test("retries bounded cleanup after a write-probe removal times out", async () => {
  const adapters = successfulAdapters({
    removeDiagnosticRecord: vi.fn(() => new Promise<void>(() => undefined)),
  });
  const diagnoseWithShortTimeout = createWorkspaceDiagnoser(adapters, { timeoutMs: 10 });

  const startedAt = Date.now();
  const result = await diagnoseWithShortTimeout(workspace, bindings, { writeProbe: true });

  expect(Date.now() - startedAt).toBeLessThan(250);
  expect(adapters.writeDiagnosticRecord).toHaveBeenCalledTimes(1);
  expect(adapters.removeDiagnosticRecord).toHaveBeenCalledTimes(2);
  expect(result).toMatchObject({ activatable: false });
  expect(JSON.stringify(result)).not.toContain("timeout");
});