Azure log analytics extension for Pi coding agent

azure-log-analytics.ts

/**
 * Azure Log Analytics Extension
 *
 * Adds the `azure_log_analytics_query` tool, which runs bounded read-only KQL
 * queries through the local Azure CLI (`az monitor log-analytics query`). The
 * tool uses the current Azure CLI login (`az login`) and that account's RBAC
 * permissions; no Azure credentials are stored by Pi.
 *
 * Prerequisites:
 * - Azure CLI is installed and available as `az`.
 * - You are logged in with `az login` and have access to the target workspace.
 * - This extension directory is enabled in Pi settings, for example:
 *   `"extensions": [".pi/global-extensions"]`.
 *
 * Configuration can be set in process environment variables, `<cwd>/.env`, or
 * `<cwd>/django/.env`:
 *
 * ```dotenv
 * # Required unless each tool call passes workspaceId explicitly.
 * AZURE_LOG_ANALYTICS_WORKSPACE_ID="<log-analytics-workspace-guid>"
 *
 * # Optional when the current `az account` is not the target subscription.
 * # AZURE_SUBSCRIPTION_ID is also accepted as a fallback.
 * AZURE_LOG_ANALYTICS_SUBSCRIPTION_ID="<subscription-id-or-name>"
 *
 * # Optional defaults.
 * AZURE_LOG_ANALYTICS_TIMESPAN="P1D"
 * AZURE_LOG_ANALYTICS_TIMEOUT_MS="120000"
 *
 * # Large query outputs are written to ./tmp/pi-azure-log-analytics/ instead of inlined.
 * AZURE_LOG_ANALYTICS_MAX_INLINE_BYTES="24000"
 * AZURE_LOG_ANALYTICS_MAX_OUTPUT_BYTES="5000000"
 * ```
 *
 * Example KQL:
 * ```kql
 * AppTraces
 * | where TimeGenerated > ago(1h)
 * | summarize count() by SeverityLevel
 * ```
 */

import { spawn } from "node:child_process";
import { existsSync, readFileSync } from "node:fs";
import { mkdir, writeFile } from "node:fs/promises";
import { join, relative } from "node:path";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

const DEFAULT_TIMESPAN = process.env.AZURE_LOG_ANALYTICS_TIMESPAN || "P1D";
const DEFAULT_MAX_INLINE_BYTES = Number(process.env.AZURE_LOG_ANALYTICS_MAX_INLINE_BYTES || 24_000);
const DEFAULT_MAX_OUTPUT_BYTES = Number(process.env.AZURE_LOG_ANALYTICS_MAX_OUTPUT_BYTES || 5_000_000);
const DEFAULT_AZ_TIMEOUT_MS = Number(process.env.AZURE_LOG_ANALYTICS_TIMEOUT_MS || 120_000);

interface ProcessResult {
	stdout: string;
	stderr: string;
	code: number | null;
	timedOut: boolean;
}

function runAz(args: string[], maxOutputBytes: number, timeoutMs: number, signal?: AbortSignal): Promise<ProcessResult> {
	return new Promise((resolve, reject) => {
		const child = spawn("az", args, {
			stdio: ["ignore", "pipe", "pipe"],
			env: process.env,
		});

		let stdout = "";
		let stderr = "";
		let settled = false;
		let timedOut = false;

		const finish = (result: ProcessResult) => {
			if (settled) return;
			settled = true;
			clearTimeout(timeout);
			resolve(result);
		};

		const kill = () => {
			if (!child.killed) child.kill("SIGTERM");
		};

		const timeout = setTimeout(() => {
			timedOut = true;
			kill();
		}, timeoutMs);

		child.stdout.on("data", (chunk) => {
			stdout += chunk.toString();
			if (Buffer.byteLength(stdout, "utf8") > maxOutputBytes) kill();
		});

		child.stderr.on("data", (chunk) => {
			stderr += chunk.toString();
			if (Buffer.byteLength(stderr, "utf8") > maxOutputBytes) kill();
		});

		child.on("error", reject);
		child.on("close", (code) => finish({ stdout, stderr, code, timedOut }));

		signal?.addEventListener("abort", kill, { once: true });
	});
}

function parseEnvValue(value: string): string {
	const trimmed = value.trim();
	if ((trimmed.startsWith('"') && trimmed.endsWith('"')) || (trimmed.startsWith("'") && trimmed.endsWith("'"))) {
		return trimmed.slice(1, -1);
	}
	return trimmed;
}

function numberEnvFallback(ctx: ExtensionContext, name: string, fallback: number): number {
	const value = envFallback(ctx, name);
	const parsed = value ? Number(value) : Number.NaN;
	return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback;
}

function envFallback(ctx: ExtensionContext, name: string): string | undefined {
	if (process.env[name]) return process.env[name];
	for (const path of [join(ctx.cwd, "django", ".env"), join(ctx.cwd, ".env")]) {
		if (!existsSync(path)) continue;
		const content = readFileSync(path, "utf8");
		for (const line of content.split(/\r?\n/)) {
			const trimmed = line.trim();
			if (!trimmed || trimmed.startsWith("#")) continue;
			const match = trimmed.match(/^([A-Za-z_][A-Za-z0-9_]*)=(.*)$/);
			if (match?.[1] === name) return parseEnvValue(match[2]);
		}
	}
	return undefined;
}

function validateQuery(query: string): string | undefined {
	const trimmed = query.trim();
	if (!trimmed) return "KQL query is empty.";
	if (/[;&`$<>]/.test(trimmed)) return "KQL query contains shell metacharacters blocked by this tool.";
	if (!/\b(take|limit|summarize|count|top|project)\b/i.test(trimmed)) {
		return "KQL query must be bounded with take, limit, summarize, count, top, or project.";
	}
	return undefined;
}

function timestampForFile(): string {
	return new Date().toISOString().replace(/[:.]/g, "-");
}

async function formatOutput(text: string, ctx: ExtensionContext, maxInlineBytes: number): Promise<{ contentText: string; details: Record<string, unknown> }> {
	const bytes = Buffer.byteLength(text, "utf8");
	if (bytes <= maxInlineBytes) return { contentText: text || "(no rows)", details: { inline: true, bytes } };

	const dir = join(ctx.cwd, "tmp", "pi-azure-log-analytics");
	await mkdir(dir, { recursive: true });
	const absolutePath = join(dir, `${timestampForFile()}-query-result.json`);
	await writeFile(absolutePath, text, "utf8");
	const relativePath = relative(ctx.cwd, absolutePath) || absolutePath;
	return {
		contentText: [
			`Query returned ${bytes} bytes, so full JSON output was written to ${relativePath}.`,
			"Use read, bash, jq, or Python against that file to inspect specific rows or summarize further.",
			"First bytes:",
			text.slice(0, 4_000),
		].join("\n"),
		details: { inline: false, bytes, absolutePath, relativePath },
	};
}

export default function (pi: ExtensionAPI) {
	pi.registerTool({
		name: "azure_log_analytics_query",
		label: "Azure Log Analytics Query",
		description: "Run a bounded read-only KQL query against Azure Log Analytics using the local Azure CLI login.",
		promptSnippet: "Run bounded read-only KQL queries against Azure Log Analytics using the user's Azure CLI identity.",
		promptGuidelines: [
			"Use azure_log_analytics_query when the user asks to inspect Azure Log Analytics logs or run KQL.",
			"Keep azure_log_analytics_query queries bounded with take, limit, top, count, summarize, project, or a narrow timespan.",
			"Prefer summarized KQL results over broad raw prod log dumps.",
		],
		parameters: Type.Object({
			query: Type.String({ description: "KQL query to run." }),
			workspaceId: Type.Optional(
				Type.String({ description: "Log Analytics workspace ID. Defaults to AZURE_LOG_ANALYTICS_WORKSPACE_ID." }),
			),
			timespan: Type.Optional(Type.String({ description: "Query timespan such as PT1H, P1D, or P7D. Defaults to P1D." })),
			subscriptionId: Type.Optional(
				Type.String({ description: "Optional Azure subscription ID/name passed to Azure CLI. Defaults to AZURE_LOG_ANALYTICS_SUBSCRIPTION_ID or AZURE_SUBSCRIPTION_ID." }),
			),
		}),
		async execute(_toolCallId, params, signal, onUpdate, ctx) {
			const workspaceId = params.workspaceId || envFallback(ctx, "AZURE_LOG_ANALYTICS_WORKSPACE_ID");
			const subscriptionId =
				params.subscriptionId || envFallback(ctx, "AZURE_LOG_ANALYTICS_SUBSCRIPTION_ID") || envFallback(ctx, "AZURE_SUBSCRIPTION_ID");
			const timespan = params.timespan || envFallback(ctx, "AZURE_LOG_ANALYTICS_TIMESPAN") || DEFAULT_TIMESPAN;
			const maxInlineBytes = numberEnvFallback(ctx, "AZURE_LOG_ANALYTICS_MAX_INLINE_BYTES", DEFAULT_MAX_INLINE_BYTES);
			const maxOutputBytes = numberEnvFallback(ctx, "AZURE_LOG_ANALYTICS_MAX_OUTPUT_BYTES", DEFAULT_MAX_OUTPUT_BYTES);
			const timeoutMs = numberEnvFallback(ctx, "AZURE_LOG_ANALYTICS_TIMEOUT_MS", DEFAULT_AZ_TIMEOUT_MS);

			if (!workspaceId) {
				return {
					isError: true,
					content: [{ type: "text", text: "Set AZURE_LOG_ANALYTICS_WORKSPACE_ID or pass workspaceId." }],
				};
			}

			const validationError = validateQuery(params.query);
			if (validationError) {
				return { isError: true, content: [{ type: "text", text: validationError }] };
			}

			onUpdate?.({ content: [{ type: "text", text: `Running Azure Log Analytics query (${timespan})...` }] });

			const args = [
				"monitor",
				"log-analytics",
				"query",
				"--workspace",
				workspaceId,
				"--analytics-query",
				params.query,
				"--timespan",
				timespan,
				"--output",
				"json",
			];
			if (subscriptionId) args.push("--subscription", subscriptionId);

			const result = await runAz(args, maxOutputBytes, timeoutMs, signal);
			if (result.timedOut) {
				return {
					isError: true,
					content: [{ type: "text", text: `Azure CLI query timed out after ${timeoutMs}ms.` }],
					details: result,
				};
			}
			if (result.code !== 0) {
				const formatted = await formatOutput(result.stderr || result.stdout || `Azure CLI exited with code ${result.code}.`, ctx, maxInlineBytes);
				return {
					isError: true,
					content: [{ type: "text", text: formatted.contentText }],
					details: { ...formatted.details, ...result, workspaceId, subscriptionId, timespan, query: params.query },
				};
			}

			const formatted = await formatOutput(result.stdout || "(no rows)", ctx, maxInlineBytes);
			return {
				content: [{ type: "text", text: formatted.contentText }],
				details: { ...formatted.details, workspaceId, subscriptionId, timespan, query: params.query },
			};
		},
	});
}