Azure log analytics extension for Pi coding agent

| /** | | | * 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 }, | | | }; | | | }, | | | }); | | | } |