detection-mentor.js

import { check, sleep } from "k6";
import crypto from "k6/crypto";
import exec from "k6/execution";
import http from "k6/http";
import { Rate, Trend } from "k6/metrics";

import {
	apiBaseUrl,
	authHeaders,
	integer,
	jsonField,
	required,
	sharedThresholds,
} from "./lib/config.js";

const workspace = required("WORKSPACE_SLUG");
const mentorVUs = integer("MENTOR_VUS", 5);
const reviewVUs = integer("REVIEW_VUS", 2);
const reviewRequests = integer("REVIEW_REQUESTS", reviewVUs);
const reviewJobDuration = new Trend("review_job_duration", true);
const reviewJobsCompleted = new Rate("review_jobs_completed");
const terminalStatuses = new Set(["COMPLETED", "FAILED", "TIMED_OUT", "CANCELLED"]);
const api = apiBaseUrl();
const headers = authHeaders();

export const options = {
	discardResponseBodies: false,
	scenarios: {
		mentor_sessions: {
			exec: "mentor",
			executor: "constant-vus",
			vus: mentorVUs,
			duration: __ENV.DURATION ?? "10m",
		},
		practice_detection: {
			exec: "detection",
			executor: "shared-iterations",
			vus: reviewVUs,
			iterations: reviewRequests,
			maxDuration: __ENV.REVIEW_MAX_DURATION ?? "20m",
			startTime: __ENV.REVIEW_START_TIME ?? "15s",
		},
	},
	thresholds: {
		...sharedThresholds,
		"checks{scenario:mentor_sessions}": ["rate>0.99"],
		"checks{scenario:practice_detection}": ["rate==1"],
		"http_req_duration{operation:mentor_turn}": ["p(95)<120000"],
		"http_req_duration{operation:review_request}": ["p(95)<1000"],
		review_job_duration: ["p(95)<900000"],
		review_jobs_completed: ["rate==1"],
	},
};

export function setup() {
	const artifactIds = required("ARTIFACT_IDS")
		.split(",")
		.map((value) => Number(value.trim()));
	if (artifactIds.some((value) => !Number.isSafeInteger(value) || value < 1))
		throw new Error("ARTIFACT_IDS must be comma-separated positive integers");
	if (reviewRequests > artifactIds.length)
		throw new Error("ARTIFACT_IDS must contain at least REVIEW_REQUESTS distinct ids");
	return { artifactIds };
}

export function mentor() {
	const messageId = crypto.randomUUID();
	const response = http.post(
		`${api}/workspaces/${encodeURIComponent(workspace)}/mentor/chat`,
		JSON.stringify({
			message: {
				id: messageId,
				role: "user",
				parts: [{ type: "text", text: "Summarize my current practice priorities." }],
			},
		}),
		{ headers, tags: { operation: "mentor_turn" }, timeout: "10m" },
	);
	check(response, {
		"mentor stream completed": (result) => result.status === 200,
		"mentor emitted completion": (result) => result.body?.includes("[DONE]") === true,
	});
}

export function detection(data) {
	const artifactId = data.artifactIds[exec.scenario.iterationInTest];
	const startedAt = Date.now();
	const response = http.post(
		`${api}/workspaces/${encodeURIComponent(workspace)}/practices/review-requests`,
		JSON.stringify({ artifactKind: "scm.pull-request", artifactId }),
		{ headers, tags: { operation: "review_request" } },
	);
	const submitted = response.status === 200 && jsonField(response, "status") === "SUBMITTED";
	const jobId = submitted ? jsonField(response, "jobId") : null;
	if (
		!check(response, { "review request submitted": () => submitted && typeof jobId === "string" })
	) {
		reviewJobsCompleted.add(false);
		return;
	}

	const deadline = startedAt + integer("REVIEW_TIMEOUT_SECONDS", 1200) * 1000;
	while (Date.now() < deadline) {
		sleep(2);
		const job = http.get(
			`${api}/workspaces/${encodeURIComponent(workspace)}/agents/jobs/${jobId}`,
			{ headers, tags: { operation: "review_status" } },
		);
		if (job.status !== 200) continue;
		const status = jsonField(job, "status");
		if (terminalStatuses.has(status)) {
			reviewJobDuration.add(Date.now() - startedAt);
			reviewJobsCompleted.add(status === "COMPLETED");
			check(job, { "review job completed": () => status === "COMPLETED" });
			return;
		}
	}
	reviewJobsCompleted.add(false);
}