pi-nebius 0.3.0

package/src/benchmark/command.ts

@@ -0,0 +1,218 @@
+import { randomUUID } from "node:crypto";
+import { mkdir, mkdtemp, readFile, realpath, rm } from "node:fs/promises";
+import { findPackageJSON } from "node:module";
+import { tmpdir } from "node:os";
+import { dirname, join } from "node:path";
+import { fileURLToPath, pathToFileURL } from "node:url";
+import { parseArgs } from "node:util";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { MISSING_KEY } from "../discovery.ts";
+import { applyModelSettings, type ModelSettingsMap } from "../model-settings.ts";
+import type { NebiusModel } from "../models.ts";
+import { loadDefinition, positive } from "./definition.ts";
+import { redactor } from "./instrumentation.ts";
+import { snapshotProject } from "./project.ts";
+import { terminalReport } from "./report.ts";
+import { runBenchmark } from "./runner.ts";
+import type { BenchmarkDefinition } from "./types.ts";
+
+const tasks = ["fix-auth-bug", "add-api-endpoint", "refactor-module", "multi-file-feature"];
+const help = `Usage: /nebius-benchmark [--models ID,ID] [--runs N]
+Enter your task prompt in the editor. Each run gets a fresh copy of the current project.
+Without --models, uses the selected Nebius model. Default: 1 run per model.
+Optional: --task NAME uses a bundled task instead (${tasks.join(", ")}).
+/nebius-benchmark cancel stops the active benchmark.
+Results appear here and in benchmark-results/. Custom prompts have no correctness check.
+Copies exclude Git-ignored files, dependencies, build output, and known credential files.
+Runs use paid inference and tools with normal host access; copies are not a security sandbox.`;
+
+export async function hostPiEntry(): Promise<string> {
+  // The running Pi CLI can be a symlink into a global installation.
+  const base = process.argv[1]
+    ? pathToFileURL(await realpath(process.argv[1])).href
+    : import.meta.url;
+  const manifestPath = findPackageJSON("@earendil-works/pi-coding-agent", base);
+  if (!manifestPath)
+    throw new Error("Cannot locate the host Pi SDK. Use the Node.js Pi installation.");
+  const manifest = JSON.parse(await readFile(manifestPath, "utf8"));
+  return join(dirname(manifestPath), manifest.main);
+}
+
+export function registerBenchmarkCommand(
+  pi: ExtensionAPI,
+  getSettings: () => ModelSettingsMap = () => ({}),
+) {
+  let starting = false;
+  let active: { controller: AbortController; done: Promise<void> } | undefined;
+  const show = (content: string) =>
+    pi.sendMessage(
+      { customType: "nebius-benchmark", content, display: true },
+      { triggerTurn: false },
+    );
+  pi.registerCommand("nebius-benchmark", {
+    description: "Compare models on your prompt: --models ID,ID --runs N; help or cancel",
+    handler: async (args, ctx) => {
+      const tokens = args.trim().split(/\s+/).filter(Boolean);
+      const action = tokens[0];
+      if (action === "help" || action === "list") {
+        show(help);
+        return;
+      }
+      if (action === "cancel") {
+        if (active) {
+          active.controller.abort();
+          ctx.ui.notify("Stopping benchmark…", "info");
+        } else ctx.ui.notify("No benchmark is running.", "info");
+        return;
+      }
+      if (active || starting) {
+        ctx.ui.notify("A benchmark is already running. Use /nebius-benchmark cancel.", "warning");
+        return;
+      }
+      const key = process.env.NEBIUS_API_KEY?.trim();
+      const redact = redactor([key ?? ""]);
+      starting = true;
+      let scratch: string | undefined;
+      try {
+        const { values } = parseArgs({
+          args: tokens,
+          options: {
+            models: { type: "string" },
+            runs: { type: "string" },
+            task: { type: "string" },
+          },
+        });
+        const runs = positive(Number(values.runs ?? 1), "runs", 1000);
+        if (values.task && !tasks.includes(values.task)) throw new Error(help);
+        let selected = ctx.model ? [ctx.model] : [];
+        if (values.models !== undefined) {
+          const ids = values.models.split(",");
+          if (ids.some((id) => !id) || new Set(ids).size !== ids.length)
+            throw new Error("--models requires unique, comma-separated model IDs.");
+          const available = ctx.modelRegistry.getAll();
+          selected = ids.map((id) => {
+            const model = available.find(
+              (candidate) => candidate.provider === "nebius" && candidate.id === id,
+            );
+            if (!model)
+              throw new Error(`Unknown Nebius model: ${id}. Run /nebius-refresh or check /model.`);
+            return model;
+          });
+        }
+        if (
+          !selected.length ||
+          selected.some(
+            (model) => model.provider !== "nebius" || model.api !== "openai-completions",
+          )
+        )
+          throw new Error("Select a Nebius model with /model first.");
+        const total = runs * selected.length;
+        if (total > 10000) throw new Error("At most 10,000 runs per benchmark.");
+        if (!key) throw new Error(MISSING_KEY);
+        const piEntry = await hostPiEntry();
+        let definition: BenchmarkDefinition;
+        let directory: string;
+        const task = values.task ?? "custom-prompt";
+        if (values.task) {
+          const root = fileURLToPath(new URL("../../", import.meta.url));
+          ({ definition, directory } = await loadDefinition(join(root, "benchmarks", task)));
+        } else {
+          if (!ctx.hasUI)
+            throw new Error("Custom prompts need interactive Pi. Use --task for a bundled task.");
+          const prompt = await ctx.ui.editor("Benchmark task — what should each model do?");
+          if (!prompt?.trim()) return;
+          scratch = await mkdtemp(join(tmpdir(), "pi-nebius-project-"));
+          directory = scratch;
+          const copied = await snapshotProject(ctx.cwd, join(directory, "fixture"));
+          await mkdir(join(directory, "validation"));
+          definition = {
+            schemaVersion: 1,
+            name: task,
+            task: prompt.trim(),
+            fixture: "fixture",
+            validationDirectory: "validation",
+            validation: [],
+            validationMode: "none",
+            setup: [],
+            timeout: 600,
+            validationTimeout: 60,
+            tools: ["read", "bash", "edit", "write", "grep", "find", "ls"],
+          };
+          ctx.ui.notify(
+            `Copied ${copied} project files for comparison. Dependencies and ignored files are excluded.`,
+            "info",
+          );
+        }
+        // Snapshot selection; changing the interactive model does not change an active run.
+        const modelSettings = structuredClone(getSettings());
+        const models = (structuredClone(selected) as NebiusModel[]).map((model) =>
+          applyModelSettings(model, modelSettings[model.id]),
+        );
+        let completed = 0;
+        const output = join(ctx.cwd, "benchmark-results", `${task}-${randomUUID()}`);
+        const controller = new AbortController();
+        show(
+          `Starting ${task}: ${models.map((model) => model.id).join(", ")}, ${runs} run(s) each. Uses paid inference; tools have normal host access.\nUse /nebius-benchmark cancel to stop.`,
+        );
+        ctx.ui.setStatus("nebius-benchmark", `Benchmark: ${task} (0/${total})`);
+        const snapshot = scratch;
+        scratch = undefined; // The background run owns cleanup from here.
+        const done = runBenchmark({
+          definition,
+          directory,
+          models,
+          modelSettings,
+          runs,
+          concurrency: 1,
+          output,
+          apiKey: key,
+          signal: controller.signal,
+          piEntry,
+          workerPath: fileURLToPath(
+            new URL(
+              import.meta.url.endsWith(".ts") ? "./host-worker.ts" : "./host-worker.js",
+              import.meta.url,
+            ),
+          ),
+          workerArgs: [piEntry],
+          onRun: (run) => {
+            ctx.ui.setStatus("nebius-benchmark", `Benchmark: ${task} (${++completed}/${total})`);
+            ctx.ui.notify(
+              `${run.model} #${run.run}: ${run.failure ?? (run.validation.checked ? "PASS" : "FINISHED (not validated)")}`,
+              "info",
+            );
+          },
+        })
+          .then((results) => {
+            show(redact(`${terminalReport(results)}\nDetails: ${join(output, "results.json")}`));
+          })
+          .catch((error) => {
+            show(redact(`Benchmark failed: ${String(error)}`));
+          })
+          .finally(async () => {
+            try {
+              if (snapshot) await rm(snapshot, { recursive: true, force: true });
+            } catch (error) {
+              ctx.ui.notify(
+                redact(`Could not remove temporary snapshot: ${String(error)}`),
+                "warning",
+              );
+            } finally {
+              active = undefined;
+              ctx.ui.setStatus("nebius-benchmark", undefined);
+            }
+          });
+        active = { controller, done };
+      } catch (error) {
+        ctx.ui.notify(redact(String(error)), "error");
+      } finally {
+        if (scratch) await rm(scratch, { recursive: true, force: true });
+        starting = false;
+      }
+    },
+  });
+  pi.on("session_shutdown", async () => {
+    active?.controller.abort();
+    await active?.done;
+  });
+}

package/src/benchmark/definition.ts

@@ -0,0 +1,112 @@
+import { readFile, realpath } from "node:fs/promises";
+import { dirname, isAbsolute, join, relative, resolve, sep } from "node:path";
+import { parse } from "yaml";
+import { isRecord } from "../models.ts";
+import type { BenchmarkDefinition, Command } from "./types.ts";
+
+export function positive(value: unknown, name: string, maximum = 86400): number {
+  if (typeof value !== "number" || !Number.isSafeInteger(value) || value < 1 || value > maximum)
+    throw new Error(`${name} must be an integer between 1 and ${maximum}`);
+  return value;
+}
+function text(value: unknown, name: string): string {
+  if (typeof value !== "string" || !value.trim()) throw new Error(`${name} must be nonempty text`);
+  return value;
+}
+function fields(value: Record<string, unknown>, allowed: string[]) {
+  for (const key of Object.keys(value))
+    if (!allowed.includes(key)) throw new Error(`Unknown field: ${key}`);
+}
+function commands(value: unknown, name: string): Command[] {
+  if (!Array.isArray(value)) throw new Error(`${name} must be an array`);
+  return value.map((item) => {
+    if (typeof item === "string") return { command: "/bin/sh", args: ["-c", text(item, name)] };
+    if (!isRecord(item)) throw new Error(`Invalid ${name} command`);
+    fields(item, ["command", "args"]);
+    if (
+      item.args !== undefined &&
+      (!Array.isArray(item.args) || item.args.some((arg) => typeof arg !== "string"))
+    )
+      throw new Error(`${name} args must be strings`);
+    return { command: text(item.command, "command"), args: (item.args ?? []) as string[] };
+  });
+}
+function localPath(value: unknown, fallback: string): string {
+  const path = value === undefined ? fallback : text(value, "path");
+  if (isAbsolute(path) || path.split(/[\\/]/).includes(".."))
+    throw new Error("Benchmark paths must be relative and inside its directory");
+  return path;
+}
+export function parseDefinition(source: string): BenchmarkDefinition {
+  const value: unknown = parse(source, { maxAliasCount: 0, uniqueKeys: true });
+  if (!isRecord(value)) throw new Error("Benchmark definition must be an object");
+  fields(value, [
+    "schemaVersion",
+    "name",
+    "task",
+    "fixture",
+    "validationDirectory",
+    "validation",
+    "setup",
+    "timeout",
+    "validationTimeout",
+    "tools",
+    "systemPrompt",
+  ]);
+  if (value.schemaVersion !== undefined && value.schemaVersion !== 1)
+    throw new Error("Unsupported benchmark schemaVersion");
+  const tools = value.tools ?? ["read", "bash", "edit", "write"];
+  const available = ["read", "bash", "edit", "write", "grep", "find", "ls"];
+  if (
+    !Array.isArray(tools) ||
+    !tools.length ||
+    tools.some((tool) => !available.includes(tool)) ||
+    new Set(tools).size !== tools.length
+  )
+    throw new Error("tools must be a nonempty, unique list of Pi built-in tools");
+  const validation = commands(value.validation, "validation");
+  if (!validation.length) throw new Error("At least one validation command is required");
+  return {
+    schemaVersion: 1,
+    name: text(value.name, "name"),
+    task: text(value.task, "task"),
+    fixture: localPath(value.fixture, "fixture"),
+    validationDirectory: localPath(value.validationDirectory, "validation"),
+    validation,
+    setup: commands(value.setup ?? [], "setup"),
+    timeout: positive(value.timeout ?? 600, "timeout"),
+    validationTimeout: positive(value.validationTimeout ?? 60, "validationTimeout"),
+    tools,
+    ...(value.systemPrompt === undefined
+      ? {}
+      : { systemPrompt: text(value.systemPrompt, "systemPrompt") }),
+  };
+}
+export async function loadDefinition(path: string) {
+  const definitionFile =
+    path.endsWith(".yaml") || path.endsWith(".yml") || path.endsWith(".json")
+      ? resolve(path)
+      : resolve(path, "benchmark.yaml");
+  const source = await readFile(definitionFile, "utf8");
+  if (Buffer.byteLength(source) > 1024 * 1024) throw new Error("Definition exceeds 1 MiB");
+  const definition = parseDefinition(source);
+  const directory = await realpath(dirname(definitionFile));
+  const directories: string[] = [];
+  for (const child of [definition.fixture, definition.validationDirectory]) {
+    const actual = await realpath(join(directory, child));
+    directories.push(actual);
+    const rel = relative(directory, actual);
+    if (!rel || rel.startsWith("..") || isAbsolute(rel))
+      throw new Error("Fixture/validation must be separate directories inside the benchmark");
+  }
+  const [fixture, validators] = directories;
+  if (
+    fixture &&
+    validators &&
+    [relative(fixture, validators), relative(validators, fixture)].some(
+      (path) => !isAbsolute(path) && path !== ".." && !path.startsWith(`..${sep}`),
+    )
+  )
+    throw new Error("Fixture and validation directories must be separate and must not overlap");
+  return { definition, directory, source };
+}

package/src/benchmark/host-worker.ts

@@ -0,0 +1,14 @@
+import { registerHooks } from "node:module";
+import { pathToFileURL } from "node:url";
+
+// Reuse the installed host's SDK; Git installs intentionally omit Pi dev dependencies.
+const hostEntry = process.argv[2];
+if (!hostEntry) throw new Error("Missing host Pi entry");
+registerHooks({
+  resolve(specifier, context, nextResolve) {
+    if (specifier.startsWith("@earendil-works/"))
+      return nextResolve(specifier, { ...context, parentURL: pathToFileURL(hostEntry).href });
+    return nextResolve(specifier, context);
+  },
+});
+await import("./worker.ts");

package/src/benchmark/instrumentation.ts

@@ -0,0 +1,298 @@
+import { createHash } from "node:crypto";
+import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent";
+import { isRecord } from "../models.ts";
+import type { Observation, ReportedUsage, RequestTrace } from "./types.ts";
+
+export const hash = (text: string) => createHash("sha256").update(text).digest("hex");
+const count = (value: unknown): number | null =>
+  typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : null;
+export function reportedUsage(value: unknown): ReportedUsage | null {
+  if (!isRecord(value)) return null;
+  const prompt = isRecord(value.prompt_tokens_details) ? value.prompt_tokens_details : {};
+  const completion = isRecord(value.completion_tokens_details)
+    ? value.completion_tokens_details
+    : {};
+  return {
+    inputTokens: count(value.prompt_tokens),
+    outputTokens: count(value.completion_tokens),
+    cachedInputTokens: count(
+      prompt.cached_tokens ?? value.prompt_cache_hit_tokens ?? value.cached_tokens,
+    ),
+    reasoningTokens: count(completion.reasoning_tokens),
+    totalTokens: count(value.total_tokens),
+  };
+}
+export function emptyObservation(): Observation {
+  return {
+    requests: [],
+    tools: [],
+    agentTurns: 0,
+    assistantMessages: 0,
+    toolCalls: 0,
+    toolCallsByType: {},
+    compactions: 0,
+    retries: 0,
+    agentStartedAtMs: null,
+    agentEndedAtMs: null,
+    systemPromptHash: null,
+    lastAssistantStopReason: null,
+    errors: [],
+  };
+}
+export function redactor(secrets: string[]) {
+  return (value: string) => {
+    let safe = value;
+    for (const secret of secrets.filter(Boolean)) safe = safe.split(secret).join("[REDACTED]");
+    return safe.replace(/Bearer\s+[^\s"']+/gi, "Bearer [REDACTED]");
+  };
+}
+
+/** Observer only: never returns modified payloads, tools, messages, or stream bytes. */
+export class Instrumentation {
+  readonly state = emptyObservation();
+  private compacting = false;
+  private readonly update: (state: Observation) => void;
+  private readonly now: () => number;
+  private readonly redact: (text: string) => string;
+  constructor(
+    update: (state: Observation) => void = () => {},
+    now: () => number = () => performance.now(),
+    redact: (text: string) => string = (text) => text,
+  ) {
+    this.update = update;
+    this.now = now;
+    this.redact = redact;
+  }
+
+  systemPrompt(prompt: string) {
+    this.state.systemPromptHash = hash(prompt);
+    this.publish();
+  }
+  private publish() {
+    this.update(this.state);
+  }
+  onEvent(event: AgentSessionEvent) {
+    const at = this.now();
+    switch (event.type) {
+      case "agent_start":
+        this.state.agentStartedAtMs ??= at;
+        break;
+      case "agent_settled":
+        this.state.agentEndedAtMs = at;
+        break;
+      case "turn_start":
+        this.state.agentTurns++;
+        break;
+      case "compaction_start":
+        this.compacting = true;
+        this.state.compactions++;
+        break;
+      case "compaction_end":
+        this.compacting = false;
+        break;
+      case "auto_retry_start":
+        this.state.retries++;
+        break;
+      case "tool_execution_start":
+        this.state.tools.push({
+          id: event.toolCallId,
+          name: event.toolName,
+          startedAtMs: at,
+          endedAtMs: null,
+          isError: null,
+          error: null,
+        });
+        break;
+      case "tool_execution_end": {
+        const tool = this.state.tools.find((item) => item.id === event.toolCallId);
+        if (tool) {
+          tool.endedAtMs = at;
+          tool.isError = event.isError;
+        }
+        break;
+      }
+      case "message_end": {
+        const message = event.message;
+        if (message.role === "assistant") {
+          this.state.assistantMessages++;
+          this.state.lastAssistantStopReason = message.stopReason;
+          if (message.errorMessage)
+            this.state.errors.push(this.redact(message.errorMessage).slice(0, 8000));
+          for (const block of message.content)
+            if (block.type === "toolCall") {
+              this.state.toolCalls++;
+              this.state.toolCallsByType[block.name] =
+                (this.state.toolCallsByType[block.name] ?? 0) + 1;
+            }
+        } else if (message.role === "toolResult") {
+          let tool = this.state.tools.find((item) => item.id === message.toolCallId);
+          if (!tool) {
+            tool = {
+              id: message.toolCallId,
+              name: message.toolName,
+              startedAtMs: null,
+              endedAtMs: at,
+              isError: message.isError,
+              error: null,
+            };
+            this.state.tools.push(tool);
+          }
+          tool.isError = message.isError;
+          if (message.isError)
+            tool.error = this.redact(
+              message.content
+                .filter((item) => item.type === "text")
+                .map((item) => item.text)
+                .join("\n"),
+            ).slice(0, 4000);
+        }
+        break;
+      }
+      default:
+        return;
+    }
+    this.publish();
+  }
+
+  observeFetch(fetcher: typeof fetch): typeof fetch {
+    return async (input, init) => {
+      const url = new URL(input instanceof Request ? input.url : input.toString());
+      if (!url.pathname.endsWith("/chat/completions")) return fetcher(input, init);
+      const trace: RequestTrace = {
+        request: this.state.requests.length + 1,
+        purpose: this.compacting ? "compaction" : "agent",
+        startedAtMs: this.now(),
+        endedAtMs: null,
+        firstContentAtMs: null,
+        status: null,
+        requestId: null,
+        servedModel: null,
+        systemFingerprint: null,
+        finishReason: null,
+        usage: null,
+        error: null,
+        streamComplete: false,
+      };
+      this.state.requests.push(trace);
+      this.publish();
+      try {
+        const response = await fetcher(input, init);
+        trace.status = response.status;
+        trace.requestId = response.headers.get("x-request-id");
+        this.publish();
+        if (!response.body) {
+          trace.endedAtMs = this.now();
+          this.publish();
+          return response;
+        }
+        const decoder = new TextDecoder();
+        let buffer = "";
+        let dropped = false;
+        const consume = (data: string) => {
+          if (data.trim() === "[DONE]") {
+            trace.streamComplete = true;
+            return;
+          }
+          let chunk: unknown;
+          try {
+            chunk = JSON.parse(data);
+          } catch {
+            return;
+          }
+          if (!isRecord(chunk)) return;
+          if (typeof chunk.model === "string") trace.servedModel = chunk.model;
+          if (typeof chunk.system_fingerprint === "string")
+            trace.systemFingerprint = chunk.system_fingerprint;
+          if (chunk.usage) trace.usage = reportedUsage(chunk.usage);
+          for (const choice of Array.isArray(chunk.choices) ? chunk.choices : []) {
+            if (!isRecord(choice)) continue;
+            if (typeof choice.finish_reason === "string") trace.finishReason = choice.finish_reason;
+            if (choice.usage) trace.usage = reportedUsage(choice.usage);
+            const delta = isRecord(choice.delta) ? choice.delta : {};
+            if (
+              trace.firstContentAtMs === null &&
+              ([delta.content, delta.reasoning_content, delta.reasoning].some(
+                (value) => typeof value === "string" && value.length > 0,
+              ) ||
+                (Array.isArray(delta.tool_calls) &&
+                  delta.tool_calls.some((call: unknown) => {
+                    if (!isRecord(call) || !isRecord(call.function)) return false;
+                    return [call.function.name, call.function.arguments].some(
+                      (value) => typeof value === "string" && value.length > 0,
+                    );
+                  })))
+            )
+              trace.firstContentAtMs = this.now();
+          }
+        };
+        const ingest = (bytes: Uint8Array) => {
+          if (dropped) return;
+          const previous = JSON.stringify([trace.usage, trace.firstContentAtMs]);
+          buffer += decoder.decode(bytes, { stream: true });
+          if (buffer.length > 1024 * 1024) {
+            dropped = true;
+            buffer = "";
+            trace.error = "Instrumentation frame exceeded 1 MiB; usage may be incomplete";
+            return;
+          }
+          // Line parser handles arbitrary TCP chunking, LF and CRLF. JSON payloads
+          // are contained in SSE data lines for Nebius's Chat Completions protocol.
+          let newline = buffer.indexOf("\n");
+          while (newline >= 0) {
+            const line = buffer.slice(0, newline).replace(/\r$/, "");
+            buffer = buffer.slice(newline + 1);
+            if (line.startsWith("data:")) consume(line.slice(5).trimStart());
+            newline = buffer.indexOf("\n");
+          }
+          if (JSON.stringify([trace.usage, trace.firstContentAtMs]) !== previous) this.publish();
+        };
+        const reader = response.body.getReader();
+        const finish = () => {
+          trace.endedAtMs = this.now();
+          this.publish();
+        };
+        const body = new ReadableStream<Uint8Array>({
+          pull: async (controller) => {
+            try {
+              const result = await reader.read();
+              if (result.done) {
+                buffer += decoder.decode();
+                if (buffer.startsWith("data:")) consume(buffer.slice(5).trim());
+                if (!response.ok)
+                  trace.error = this.redact(`HTTP ${response.status}: ${buffer}`).slice(0, 8000);
+                finish();
+                controller.close();
+                return;
+              }
+              // Forward exactly the original bytes. Observation is best effort and cannot fail inference.
+              try {
+                ingest(result.value);
+              } catch {
+                trace.error = "Instrumentation could not decode response metadata";
+              }
+              controller.enqueue(result.value);
+            } catch (error) {
+              trace.error = this.redact(String(error)).slice(0, 8000);
+              finish();
+              controller.error(error);
+            }
+          },
+          cancel: async (reason) => {
+            finish();
+            await reader.cancel(reason);
+          },
+        });
+        return new Response(body, {
+          status: response.status,
+          statusText: response.statusText,
+          headers: response.headers,
+        });
+      } catch (error) {
+        trace.error = this.redact(String(error)).slice(0, 8000);
+        trace.endedAtMs = this.now();
+        this.publish();
+        throw error;
+      }
+    };
+  }
+}