@gentrace/core
Version:
Core Gentrace Node.JS library
551 lines (548 loc) • 20 kB
JavaScript
import { GENTRACE_API_KEY, globalGentraceConfig, GENTRACE_ENVIRONMENT_NAME, getGentraceBasePath } from './init.mjs';
import { Pipeline } from './pipeline.mjs';
import { updateTestResultWithRunners } from './runners.mjs';
import WebSocket from 'ws';
import { AsyncLocalStorage } from 'async_hooks';
import * as Mustache from 'mustache';
import { BASE_PATH } from '../base.mjs';
var __awaiter = (undefined && undefined.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
const interactions = {};
function defineInteraction(interaction) {
interactions[interaction.name] = interaction;
Object.values(listeners).forEach((listener) => listener({
type: "register-interaction",
interaction,
}));
return interaction.fn;
}
const testSuites = {};
function defineTestSuite(testSuite) {
testSuites[testSuite.name] = testSuite;
Object.values(listeners).forEach((listener) => listener({
type: "register-test-suite",
testSuite,
}));
return testSuite.fn;
}
const getWSBasePath = () => {
const apiBasePath = getGentraceBasePath();
if (apiBasePath === "") {
return "wss://gentrace.ai/ws";
}
if (apiBasePath.includes("localhost")) {
return "ws://localhost:3001";
}
return ("wss://" +
apiBasePath.slice(apiBasePath.indexOf("/") + 2, apiBasePath.lastIndexOf("/")) +
"/ws");
};
const listeners = {};
const makeUuid = () => {
// Generate 16 random bytes
const bytes = new Array(16);
for (let i = 0; i < 16; i++) {
bytes[i] = Math.floor(Math.random() * 256);
}
// Set the version number to 4 (UUID version 4)
bytes[6] = (bytes[6] & 0x0f) | 0x40;
// Set the variant to 10xxxxxx (RFC 4122 variant)
bytes[8] = (bytes[8] & 0x3f) | 0x80;
// Convert bytes to hex and format as UUID
const hexBytes = bytes.map((byte) => ("0" + byte.toString(16)).slice(-2));
return [
hexBytes.slice(0, 4).join(""),
hexBytes.slice(4, 6).join(""),
hexBytes.slice(6, 8).join(""),
hexBytes.slice(8, 10).join(""),
hexBytes.slice(10, 16).join(""),
].join("-");
};
const validate = (interaction, { id, inputs }) => __awaiter(void 0, void 0, void 0, function* () {
var _a;
if (!interaction.inputType) {
return {
status: "failure",
error: "No input validator found",
id,
};
}
try {
yield interaction.inputType.parseAsync(inputs);
return { status: "success", id };
}
catch (e) {
return {
status: "failure",
error: (_a = e === null || e === void 0 ? void 0 : e.message) !== null && _a !== void 0 ? _a : "Validation failed",
id,
};
}
});
const overridesAsyncLocalStorage = new AsyncLocalStorage();
const makeGetValue = (name, defaultValue) => {
return () => {
var _a;
const overrides = overridesAsyncLocalStorage.getStore();
return (_a = overrides === null || overrides === void 0 ? void 0 : overrides[name]) !== null && _a !== void 0 ? _a : defaultValue;
};
};
const parameters = {};
const numericParameter = ({ name, defaultValue, }) => {
parameters[name] = {
name,
type: "numeric",
defaultValue,
};
return {
name,
getValue: makeGetValue(name, defaultValue),
};
};
const stringParameter = ({ name, defaultValue, }) => {
parameters[name] = {
name,
type: "string",
defaultValue,
};
return {
name,
getValue: makeGetValue(name, defaultValue),
};
};
const enumParameter = ({ name, options, defaultValue, }) => {
parameters[name] = {
name,
type: "enum",
defaultValue,
options,
};
return {
name,
getValue: makeGetValue(name, defaultValue),
};
};
const templateParameter = ({ name, defaultValue, variables, }) => {
parameters[name] = {
name,
type: "template",
defaultValue,
variables: variables ? variables : [],
};
return {
name,
render: (values) => {
var _a;
const overrides = overridesAsyncLocalStorage.getStore();
const template = (_a = overrides === null || overrides === void 0 ? void 0 : overrides[name]) !== null && _a !== void 0 ? _a : defaultValue;
return Mustache.render(template, values);
},
};
};
function makeParallelRunner(parallelism) {
const results = [];
const queue = [];
let numRunning = 0;
function processQueue() {
while ((!parallelism || numRunning < parallelism) && queue.length > 0) {
const { fn, resolve, reject } = queue.shift();
numRunning++;
fn()
.then(resolve)
.catch(reject)
.finally(() => {
numRunning--;
processQueue();
});
}
}
return {
results,
run: (fn) => __awaiter(this, void 0, void 0, function* () {
results.push(new Promise((res, rej) => {
queue.push({ fn, resolve: res, reject: rej });
}));
processQueue();
}),
};
}
function sendMessage(message, transport) {
return __awaiter(this, void 0, void 0, function* () {
if (transport.type === "ws") {
if (!transport.pluginId) {
transport.messageQueue.push(message);
return;
}
if (transport.isClosed) {
return;
}
transport.ws.send(JSON.stringify({
id: makeUuid(),
for: transport.pluginId,
data: message,
}));
}
else {
transport.sendResponse(message);
}
});
}
const handleRunInteractionInputValidation = (message, transport) => __awaiter(void 0, void 0, void 0, function* () {
const { id, interactionName, data: testCases } = message;
const interaction = interactions[interactionName];
if (!interaction) {
sendMessage({
type: "run-interaction-input-validation-results",
id,
interactionName,
data: testCases.map((tc) => ({
id: tc.id,
status: "failure",
error: `Interaction ${interactionName} not found`,
})),
}, transport);
}
const validationResults = yield Promise.all(testCases.map((testCase) => validate(interaction, testCase)));
sendMessage({
type: "run-interaction-input-validation-results",
id,
interactionName,
data: validationResults,
}, transport);
});
const runTestCaseThroughInteraction = (pipelineId, testJobId, interactionName, testCase) => __awaiter(void 0, void 0, void 0, function* () {
const pipeline = new Pipeline({
id: pipelineId,
});
const interaction = interactions[interactionName];
if (!interaction) {
// TODO: submit error to gentrace
return;
}
const runner = pipeline.start();
try {
try {
yield runner.measure(interaction.fn, [testCase.inputs]);
}
catch (e) {
runner.setError(e.toString());
}
yield updateTestResultWithRunners(testJobId, [
[runner, { id: testCase.id }],
]);
}
catch (e) {
// TODO: submit error to gentrace
console.error(e);
throw e;
}
});
const handleRunTestInteraction = (message) => __awaiter(void 0, void 0, void 0, function* () {
const { testJobId, pipelineId, interactionName, parallelism, data: testCases, overrides, } = message;
const interaction = interactions[interactionName];
if (!interaction) {
return;
}
overridesAsyncLocalStorage.run(overrides, () => __awaiter(void 0, void 0, void 0, function* () {
var _b, _c;
const overrides = overridesAsyncLocalStorage.getStore();
console.log("overrides", overrides);
const parallelRunner = makeParallelRunner(parallelism);
for (const testCase of testCases) {
parallelRunner.run(() => runTestCaseThroughInteraction(pipelineId, testJobId, interactionName, testCase));
}
const results = yield Promise.allSettled(parallelRunner.results);
const erroredResults = results.filter((r) => r.status === "rejected");
if (erroredResults.length > 0) {
console.error("Errors in test job:", erroredResults.map((r) => r.reason));
}
const apiBasePath = globalGentraceConfig.basePath || BASE_PATH;
yield fetch(`${apiBasePath}/v1/test-result/status`, {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, ((_c = (_b = globalGentraceConfig.baseOptions) === null || _b === void 0 ? void 0 : _b.headers) !== null && _c !== void 0 ? _c : {})),
body: JSON.stringify({
id: testJobId,
finished: true,
}),
});
}));
});
const handleRunTestSuite = (message) => __awaiter(void 0, void 0, void 0, function* () {
var _d, _e;
const { testSuiteName, testJobId, pipelineId } = message;
const testSuite = testSuites[testSuiteName];
if (!testSuite) {
return;
}
yield testSuite.fn();
const apiBasePath = globalGentraceConfig.basePath || BASE_PATH;
yield fetch(`${apiBasePath}/v1/test-result/status`, {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, ((_e = (_d = globalGentraceConfig.baseOptions) === null || _d === void 0 ? void 0 : _d.headers) !== null && _e !== void 0 ? _e : {})),
body: JSON.stringify({
id: testJobId,
finished: true,
}),
});
});
const onInteraction = (interaction, transport) => {
var _a, _b;
sendMessage({
type: "register-interaction",
interaction: {
name: interaction.name,
hasValidation: !!interaction.inputType,
parameters: (_b = (_a = interaction.parameters) === null || _a === void 0 ? void 0 : _a.map(({ name }) => parameters[name]).filter((v) => !!v)) !== null && _b !== void 0 ? _b : [],
},
}, transport);
};
const onTestSuite = (testSuite, transport) => {
sendMessage({
type: "register-test-suite",
testSuite: {
name: testSuite.name,
},
}, transport);
};
const handleEnvironmentDetails = (message, transport) => __awaiter(void 0, void 0, void 0, function* () {
sendMessage({
type: "environment-details",
interactions: Object.values(interactions).map((interaction) => {
var _a, _b;
return ({
name: interaction.name,
hasValidation: !!interaction.inputType,
parameters: (_b = (_a = interaction.parameters) === null || _a === void 0 ? void 0 : _a.map(({ name }) => parameters[name])) !== null && _b !== void 0 ? _b : [],
});
}),
testSuites: Object.values(testSuites).map((testSuite) => ({
name: testSuite.name,
})),
}, transport);
});
const onMessage = (message, transport) => __awaiter(void 0, void 0, void 0, function* () {
switch (message.type) {
case "environment-details":
yield handleEnvironmentDetails(message, transport);
break;
case "run-interaction-input-validation":
yield handleRunInteractionInputValidation(message, transport);
break;
case "run-test-interaction":
// Immediately send confirmation to avoid timeout, then run everything
// else async
sendMessage({
type: "confirmation",
ok: true,
}, transport);
yield handleRunTestInteraction(message);
break;
case "run-test-suite":
sendMessage({
type: "confirmation",
ok: true,
}, transport);
yield handleRunTestSuite(message);
break;
}
});
function runWebSocket(environmentName, resolve, reject) {
return __awaiter(this, void 0, void 0, function* () {
const wsBasePath = getWSBasePath();
let env = environmentName !== null && environmentName !== void 0 ? environmentName : GENTRACE_ENVIRONMENT_NAME;
if (!env) {
try {
const os = yield import('os');
env = os.hostname();
}
catch (error) {
reject(new Error("Gentrace environment name is not set"));
return;
}
}
const transport = {
type: "ws",
pluginId: undefined,
ws: new WebSocket(wsBasePath),
isClosed: false,
messageQueue: [],
};
const id = makeUuid();
let intervals = [];
const sendMessage = (message) => {
if (!transport.pluginId) {
transport.messageQueue.push(message);
return;
}
if (transport.isClosed) {
return;
}
console.log("WebSocket sending message:", JSON.stringify(message, null, 2));
transport.ws.send(JSON.stringify({
id: makeUuid(),
for: transport.pluginId,
data: message,
}));
};
const setup = () => {
transport.messageQueue.forEach(sendMessage);
intervals.push(setInterval(() => {
console.log("sending heartbeat");
sendMessage({
type: "heartbeat",
});
}, 30 * 1000));
};
const cleanup = () => {
transport.isClosed = true;
delete listeners[id];
intervals.forEach(clearInterval);
intervals = [];
};
transport.ws.onopen = () => {
if (transport.isClosed) {
return;
}
console.log("WebSocket connection opened, sending setup message");
transport.ws.send(JSON.stringify({
id: makeUuid(),
init: "test-job-runner",
data: {
type: "setup",
environmentName: env,
apiKey: GENTRACE_API_KEY,
},
}));
// WSS layer, not plugin layer
intervals.push(setInterval(() => {
if (transport.isClosed) {
return;
}
console.log("sending ping");
transport.ws.send(JSON.stringify({
id: makeUuid(),
ping: true,
}));
}, 30 * 1000));
};
transport.ws.onmessage = (event) => __awaiter(this, void 0, void 0, function* () {
var _a;
if (transport.isClosed) {
return;
}
console.log("WebSocket message received:", event.data);
const messageWrapper = JSON.parse(typeof event.data === "string" ? event.data : event.data.toString());
if ((_a = messageWrapper === null || messageWrapper === void 0 ? void 0 : messageWrapper.data) === null || _a === void 0 ? void 0 : _a.pluginId) {
transport.pluginId = messageWrapper.data.pluginId;
setup();
return;
}
if (messageWrapper === null || messageWrapper === void 0 ? void 0 : messageWrapper.error) {
console.error("WebSocket error:", messageWrapper.error);
reject(new Error(JSON.stringify(messageWrapper.error, null, 2)));
return;
}
try {
const message = messageWrapper.data;
yield onMessage(message, transport);
}
catch (e) {
console.error("Error in WebSocket message handler:", e);
reject(e);
}
});
transport.ws.onclose = () => {
if (transport.isClosed) {
return;
}
cleanup();
console.log("WebSocket connection closed");
reject(new Error("WebSocket connection closed"));
};
transport.ws.onerror = (error) => {
console.error("Gentrace websocket error:", error);
};
// wait for close signal from process
process.on("SIGINT", () => {
if (transport.isClosed) {
return;
}
console.log("Received SIGINT, closing WebSocket connection");
cleanup();
transport.ws.close();
resolve();
});
process.on("SIGTERM", () => {
if (transport.isClosed) {
return;
}
console.log("Received SIGTERM, closing WebSocket connection");
cleanup();
transport.ws.close();
resolve();
});
Object.values(interactions).forEach((interaction) => onInteraction(interaction, transport));
Object.values(testSuites).forEach((testSuite) => onTestSuite(testSuite, transport));
listeners[id] = (event) => {
switch (event.type) {
case "register-interaction":
onInteraction(event.interaction, transport);
break;
case "register-test-suite":
onTestSuite(event.testSuite, transport);
break;
}
};
});
}
function listenInner({ environmentName, retries = 0, }) {
return __awaiter(this, void 0, void 0, function* () {
if (!GENTRACE_API_KEY) {
throw new Error("Gentrace API key is not set");
}
let isClosingProcess = false;
process.on("SIGINT", () => {
isClosingProcess = true;
});
process.on("SIGTERM", () => {
isClosingProcess = true;
});
try {
const closePromise = new Promise((resolve, reject) => runWebSocket(environmentName, resolve, reject));
yield closePromise;
}
catch (e) {
console.error("Error in WebSocket connection:", e);
if (isClosingProcess) {
return;
}
yield new Promise((resolve) => setTimeout(resolve, Math.min(Math.pow(2, retries) * 250, 10 * 1000)));
if (isClosingProcess) {
return;
}
return yield listenInner({ environmentName, retries: retries + 1 });
}
});
}
function listen(values) {
const { environmentName } = values !== null && values !== void 0 ? values : {};
return listenInner({ environmentName, retries: 0 });
}
function handleWebhook(body, sendResponse) {
return __awaiter(this, void 0, void 0, function* () {
console.log("Gentrace HTTP message received:", body);
yield onMessage(body, {
type: "http",
sendResponse,
});
});
}
export { defineInteraction, defineTestSuite, enumParameter, handleWebhook, listen, numericParameter, stringParameter, templateParameter };
//# sourceMappingURL=test-job-runner.mjs.map