ThothII/backend/test/model-completer.test.ts

import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, expect, test } from "vitest";
import {
  ModelCompletionProviderError,
  PythonModelCompleter,
} from "../src/catalog/model-completer.js";

const roots: string[] = [];

afterEach(() => {
  for (const root of roots.splice(0)) rmSync(root, { recursive: true, force: true });
});

function helper(source: string, options: { terminationGraceMs?: number } = {}) {
  const root = mkdtempSync(join(tmpdir(), "thothii-model-completer-"));
  roots.push(root);
  writeFileSync(join(root, "fake_completion_helper.py"), source, "utf8");
  return new PythonModelCompleter({
    pythonExecutable: "python3",
    cwd: root,
    helperModule: "fake_completion_helper",
    timeoutMs: 5_000,
    ...options,
  });
}

async function waitUntil(predicate: () => boolean, timeoutMs = 2_000): Promise<void> {
  const deadline = Date.now() + timeoutMs;
  while (!predicate()) {
    if (Date.now() >= deadline) throw new Error("condition was not met before timeout");
    await new Promise((resolve) => setTimeout(resolve, 10));
  }
}

test("uses the short-lived Python helper stdin/stdout protocol without process arguments", async () => {
  const completer = helper(`
import json
import pathlib
import sys

request = json.loads(sys.stdin.read())
pathlib.Path("request.json").write_text(
    json.dumps({"request": request, "argv": sys.argv}, sort_keys=True),
    encoding="utf-8",
)
sys.stdout.write(json.dumps({"ok": True, "content": "Descrizione italiana"}))
`);

  const content = await completer.complete({
    model: {
      id: "openai-mini",
      provider: "openai",
      model: "gpt-4.1-mini",
      endpoint: { baseUrl: "https://models.example.test/v1", apiVersion: "2026-08-01" },
      apiKeyEnv: "OPENAI_API_KEY",
      apiKey: "test-provider-secret",
    },
    messages: [
      { role: "system", content: "Return one description." },
      { role: "user", content: "Private metadata prompt." },
    ],
    signal: new AbortController().signal,
  });

  expect(content).toBe("Descrizione italiana");
  const captured = JSON.parse(readFileSync(join(roots[0]!, "request.json"), "utf8"));
  expect(captured.request).toEqual({
    model: "openai/gpt-4.1-mini",
    api_key: "test-provider-secret",
    messages: [
      { role: "system", content: "Return one description." },
      { role: "user", content: "Private metadata prompt." },
    ],
    api_base: "https://models.example.test/v1",
    api_version: "2026-08-01",
  });
  expect(JSON.stringify(captured.argv)).not.toMatch(/test-provider-secret|Private metadata prompt/);
});

test("omits api_key for an explicitly configured keyless endpoint", async () => {
  const completer = helper(`
import json
import pathlib
import sys

request = json.loads(sys.stdin.read())
pathlib.Path("request.json").write_text(json.dumps(request, sort_keys=True), encoding="utf-8")
sys.stdout.write(json.dumps({"ok": True, "content": "Descrizione Qwen"}))
`);

  await expect(completer.complete({
    model: {
      id: "qwen-36",
      provider: "openai",
      model: "qwen3.6-35b-a3b",
      disableThinking: true,
      endpoint: { baseUrl: "https://models.internal.example/v1" },
    },
    messages: [{ role: "user", content: "Describe invented metadata." }],
    signal: new AbortController().signal,
  })).resolves.toBe("Descrizione Qwen");

  expect(JSON.parse(readFileSync(join(roots[0]!, "request.json"), "utf8"))).toEqual({
    model: "openai/qwen3.6-35b-a3b",
    messages: [{ role: "user", content: "Describe invented metadata." }],
    api_base: "https://models.internal.example/v1",
    disable_thinking: true,
  });
});

test("normalizes helper failures and rejects non-pristine stdout without leaking diagnostics", async () => {
  const secret = "test-provider-secret";
  const prompt = "private metadata prompt";
  const completers = [
    helper(`
import json
import sys
request = json.loads(sys.stdin.read())
print(request["api_key"] + " " + request["messages"][0]["content"], file=sys.stderr)
sys.stdout.write(json.dumps({"ok": False, "error": "provider_failure"}))
`),
    helper(`
import json
import sys
sys.stdin.read()
sys.stdout.write(json.dumps({"ok": True, "content": "first"}) + "\\n" + json.dumps({"ok": True, "content": "second"}))
`),
  ];

  for (const completer of completers) {
    let failure: unknown;
    try {
      await completer.complete({
        model: {
          id: "openai-mini",
          provider: "openai",
          model: "gpt-4.1-mini",
          apiKeyEnv: "OPENAI_API_KEY",
          apiKey: secret,
        },
        messages: [{ role: "user", content: prompt }],
        signal: new AbortController().signal,
      });
    } catch (error) {
      failure = error;
    }
    expect(failure).toBeInstanceOf(ModelCompletionProviderError);
    expect(String(failure)).not.toMatch(new RegExp(`${secret}|${prompt}`));
  }
});

test("aborting a completion terminates its Python helper and returns a cancellation error", async () => {
  const completer = helper(`
import os
import pathlib
import signal
import sys
import time

sys.stdin.read()

def terminate(_signum, _frame):
    pathlib.Path("terminated.txt").write_text("SIGTERM", encoding="utf-8")
    raise SystemExit(0)

signal.signal(signal.SIGTERM, terminate)
pathlib.Path("pid.txt").write_text(str(os.getpid()), encoding="utf-8")
while True:
    time.sleep(0.05)
`);
  const controller = new AbortController();
  const completion = completer.complete({
    model: {
      id: "openai-mini",
      provider: "openai",
      model: "gpt-4.1-mini",
      apiKeyEnv: "OPENAI_API_KEY",
      apiKey: "test-provider-secret",
    },
    messages: [{ role: "user", content: "Private metadata prompt." }],
    signal: controller.signal,
  });
  const observed = completion.then(
    () => undefined,
    (error: unknown) => error,
  );
  const root = roots[0]!;
  await waitUntil(() => existsSync(join(root, "pid.txt")));
  const pid = Number(readFileSync(join(root, "pid.txt"), "utf8"));

  controller.abort();

  await expect(observed).resolves.toMatchObject({ name: "ModelCompletionCancelledError" });
  await waitUntil(() => {
    try {
      process.kill(pid, 0);
      return false;
    } catch {
      return true;
    }
  });
  expect(readFileSync(join(root, "terminated.txt"), "utf8")).toBe("SIGTERM");
});

test("aborting escalates to SIGKILL when the Python helper does not exit after SIGTERM", async () => {
  const completer = helper(`
import os
import pathlib
import signal
import sys
import time

sys.stdin.read()

def ignore_term(_signum, _frame):
    pathlib.Path("sigterm.txt").write_text("received", encoding="utf-8")

signal.signal(signal.SIGTERM, ignore_term)
pathlib.Path("pid.txt").write_text(str(os.getpid()), encoding="utf-8")
while True:
    time.sleep(0.05)
`, { terminationGraceMs: 25 });
  const controller = new AbortController();
  const observed = completer.complete({
    model: {
      id: "openai-mini",
      provider: "openai",
      model: "gpt-4.1-mini",
      apiKeyEnv: "OPENAI_API_KEY",
      apiKey: "test-provider-secret",
    },
    messages: [{ role: "user", content: "Private metadata prompt." }],
    signal: controller.signal,
  }).then(
    () => undefined,
    (error: unknown) => error,
  );
  const root = roots[0]!;
  await waitUntil(() => existsSync(join(root, "pid.txt")));
  const pid = Number(readFileSync(join(root, "pid.txt"), "utf8"));

  controller.abort();

  await expect(observed).resolves.toMatchObject({ name: "ModelCompletionCancelledError" });
  expect(readFileSync(join(root, "sigterm.txt"), "utf8")).toBe("received");
  await waitUntil(() => {
    try {
      process.kill(pid, 0);
      return false;
    } catch {
      return true;
    }
  });
});