import { spawn } from "node:child_process";
import { z } from "zod";
const CliMessageSchema = z.object({
role: z.enum(["system", "user", "assistant"]),
content: z.string(),
});
const CliRequestSchema = z.object({
version: z.literal("v1"),
model: z.string().min(1),
temperature: z.number(),
maxOutputTokens: z.number().int().positive(),
messages: z.array(CliMessageSchema).min(1),
metadata: z
.object({
queryId: z.string().min(1),
mode: z.string().min(1),
runKind: z.enum(["cold", "warm"]),
tokenBudget: z.number().int().positive(),
projectPath: z.string().min(1).optional(),
})
.optional(),
});
const CliResponseSchema = z.object({
outputText: z.string(),
usage: z.object({
inputTokens: z.number().int().nonnegative(),
outputTokens: z.number().int().nonnegative(),
}),
});
export type LlmCliRequest = z.infer<typeof CliRequestSchema>;
export interface LlmCliResult {
outputText: string;
inputTokens: number;
outputTokens: number;
}
export interface RunCliLlmInput {
command: string;
timeoutMs: number;
request: LlmCliRequest;
}
interface CliProcessResult {
status: number | null;
signal: NodeJS.Signals | null;
stdout: string;
stderr: string;
timedOut: boolean;
}
function parseJsonFromStdout(stdout: string): unknown {
const trimmed = stdout.trim();
if (!trimmed) {
throw new Error("CLI returned empty stdout.");
}
try {
return JSON.parse(trimmed) as unknown;
} catch {
const lines = trimmed
.split("\n")
.map((line) => line.trim())
.filter(Boolean);
for (let index = lines.length - 1; index >= 0; index--) {
try {
return JSON.parse(lines[index] ?? "") as unknown;
} catch {
// Continue scanning possible trailing JSON lines.
}
}
}
throw new Error("CLI stdout did not contain valid JSON.");
}
async function runCliProcess(input: {
command: string;
timeoutMs: number;
payload: string;
}): Promise<CliProcessResult> {
return await new Promise((resolve, reject) => {
const child = spawn(input.command, {
shell: true,
stdio: ["pipe", "pipe", "pipe"],
});
let stdout = "";
let stderr = "";
let timedOut = false;
let settled = false;
child.stdout.setEncoding("utf-8");
child.stderr.setEncoding("utf-8");
child.stdout.on("data", (chunk: string) => {
stdout += chunk;
});
child.stderr.on("data", (chunk: string) => {
stderr += chunk;
});
const timeoutMs = Math.max(1, input.timeoutMs);
const timeoutId = setTimeout(() => {
timedOut = true;
child.kill("SIGKILL");
}, timeoutMs);
child.on("error", (error) => {
if (settled) return;
settled = true;
clearTimeout(timeoutId);
reject(error);
});
child.on("close", (status, signal) => {
if (settled) return;
settled = true;
clearTimeout(timeoutId);
resolve({
status,
signal,
stdout,
stderr,
timedOut,
});
});
child.stdin.write(input.payload);
child.stdin.end();
});
}
export async function runLlmCliCompletion(
input: RunCliLlmInput,
): Promise<LlmCliResult> {
const command = input.command.trim();
if (!command) {
throw new Error(
"DOCLEA_LIVE_LLM_CLI_COMMAND is required for measured timing mode.",
);
}
const request = CliRequestSchema.parse(input.request);
const requestPayload = `${JSON.stringify(request)}\n`;
const result = await runCliProcess({
command,
timeoutMs: input.timeoutMs,
payload: requestPayload,
});
if (result.timedOut) {
throw new Error(
`CLI execution timed out after ${Math.max(1, input.timeoutMs)}ms.\nstdout:\n${result.stdout}\nstderr:\n${result.stderr}`,
);
}
if (result.status !== 0) {
throw new Error(
`CLI exited with status ${result.status}${result.signal ? ` (signal ${result.signal})` : ""}.\nstdout:\n${result.stdout}\nstderr:\n${result.stderr}`,
);
}
const parsed = parseJsonFromStdout(result.stdout);
const response = CliResponseSchema.parse(parsed);
return {
outputText: response.outputText,
inputTokens: response.usage.inputTokens,
outputTokens: response.usage.outputTokens,
};
}
