@niroai/niro 0.17.0 → 0.18.0

package/src/lib/docsUpload.js

@@ -0,0 +1,481 @@
+/**
+ * Sends the local engine's documentation requests to Niro and hands the answers back.
+ *
+ * A developer who builds a folder with `niro-local build --llm` asks for each method to be
+ * described in plain words, so `find_reusable_code` and `find_known_issues` can answer from those
+ * descriptions. The engine writes the requests, but it holds no account and no key, by design:
+ * anything that leaves the machine leaves through this CLI, which is the part that signed in.
+ *
+ * One round, the same shape as the audit and journal uploads:
+ *
+ *   niro-local docs --emit --all [--limit N]   JSON Lines, one batch of up to 12 methods each
+ *   POST <assistant>/api/mcp/local-docs/enrich  {lease, request}, per line, with the login token
+ *   niro-local docs --import                    one answer line per batch that was tried
+ *
+ * The engine's half of the contract is niro-local `docs/usecases/llm-documentation.md`, "The
+ * contract for niro-cli". ai-assistant's half is the shared contract for this change: it checks
+ * the signed licence lease sent with each request to decide whether this account's plan includes
+ * documentation, so the lease travels with every call. The lease is fetched once per round and
+ * held in memory only, as `lease.js` requires.
+ *
+ * **What never leaves this file in a log or an error.** The request carries the customer's code,
+ * the answer describes it, and the lease and the token are credentials. Log lines and failure
+ * lines carry counts, the first 12 characters of a batch id, and an HTTP status, nothing else.
+ *
+ * **Three pieces of state on disk, all under ~/.niro:**
+ *   - `docs-backoff.json`: batch hashes and times, so a batch the service keeps failing on is
+ *     asked for less and less often instead of on every round. Never code.
+ *   - `docs-send.lock`: who is running a round now, so the watcher and a person typing
+ *     `niro docs-send` never send the same batch twice at once.
+ *   - nothing else: the lease is never written down.
+ */
+const crypto = require("crypto");
+
+const api = require("./api");
+const color = require("./color");
+const lease = require("./lease");
+const localEngine = require("./localEngine");
+const periodicClaim = require("./periodicClaim");
+const { assistantUrl } = require("./journalUpload");
+const { ensureModel } = require("./modelFetch");
+
+/** Said when LLM documentation is off in this release; the engine says the same for `niro-local docs`. */
+const NOT_IN_THIS_RELEASE = "LLM documentation is not available in this release.";
+const { markedFolders, foldersMarked } = require("./docsMarked");
+const {
+  acquireLock,
+  backoffFile,
+  lockFile,
+  loadBackoff,
+  saveBackoff,
+  pruneBackoff,
+  waitingUntil,
+  recordFailure,
+  BACKOFF_BASE_MS,
+} = require("./docsState");
+
+const ENRICH_PATH = "/api/mcp/local-docs/enrich";
+
+/** How many request lines one round sends. Each line is up to 12 methods. */
+const DEFAULT_LIMIT = 20;
+
+/**
+ * How long one documentation call may take. The gateway asks an LLM about up to 12 methods in
+ * one call, which takes far longer than an upload, so this is not the 30 s the uploads use.
+ */
+const DEFAULT_TIMEOUT_MS = 180_000;
+
+/** A lease this close to its expiry is replaced before the next call. */
+const LEASE_RENEW_MARGIN_MS = 60_000;
+
+/**
+ * Read when used, not when this file is loaded, so a test or a person that sets it after
+ * startup is believed (the trap `lease.js` describes at fetchTimeoutMs).
+ */
+function timeoutMs() {
+  return Number(process.env.NIRO_DOCS_TIMEOUT_MS) || DEFAULT_TIMEOUT_MS;
+}
+
+// ---------------------------------------------------------------------------------------------
+// One round
+// ---------------------------------------------------------------------------------------------
+
+function shortId(batch) {
+  return String(batch || "?").slice(0, 12);
+}
+
+/** One line of text from a service's answer, made safe to put in a log. */
+function reasonOf(body, status) {
+  if (body && typeof body === "object") {
+    for (const field of ["error", "message", "code"]) {
+      if (typeof body[field] === "string" && body[field].trim()) {
+        return body[field].replace(/\s+/g, " ").trim().slice(0, 200);
+      }
+    }
+  }
+  return `HTTP ${status}`;
+}
+
+/** How many of this batch's methods an answer actually describes. */
+function answeredCount(body, entities) {
+  const data = body && typeof body === "object" ? body.data || body : null;
+  const results = data && Array.isArray(data.results) ? data.results : [];
+  const wanted = new Set(Object.keys(entities || {}));
+  let count = 0;
+  for (const item of results) {
+    if (
+      item &&
+      wanted.has(String(item.entityId)) &&
+      typeof item.functionality === "string" &&
+      item.functionality.trim() &&
+      typeof item.reuseDescription === "string" &&
+      item.reuseDescription.trim()
+    ) {
+      count += 1;
+    }
+  }
+  return count;
+}
+
+async function getLease(fetchLeaseImpl) {
+  const nonce = crypto.randomBytes(8).toString("hex");
+  return fetchLeaseImpl(nonce);
+}
+
+function leaseNearExpiry(held, now) {
+  const claims = lease.claimsOf(held.lease);
+  const expiresAt = claims && Number(claims.expires_at);
+  if (!Number.isFinite(expiresAt) || expiresAt <= 0) return false;
+  return expiresAt * 1000 - now < LEASE_RENEW_MARGIN_MS;
+}
+
+async function postOne(line, held, token, fetchImpl) {
+  const abort = new AbortController();
+  const timer = setTimeout(() => abort.abort(), timeoutMs());
+  try {
+    const response = await fetchImpl(`${assistantUrl()}${ENRICH_PATH}`, {
+      method: "POST",
+      headers: { "Content-Type": "application/json", Authorization: `Bearer ${token}` },
+      body: JSON.stringify({
+        lease: { payload: held.lease, signature: held.signature, key_id: held.keyId },
+        request: line.request,
+      }),
+      signal: abort.signal,
+    });
+    let body = null;
+    let unreadable = false;
+    try {
+      body = await response.json();
+    } catch {
+      unreadable = true;
+    }
+    return { status: response.status, body, unreadable };
+  } catch {
+    // The error's own message is not used: it can name the URL, and nothing else in it helps.
+    return { error: abort.signal.aborted ? "timeout" : "network error" };
+  } finally {
+    clearTimeout(timer);
+  }
+}
+
+function envelope(line) {
+  return { v: line.v, batch: line.batch, worktrees: line.worktrees, entities: line.entities };
+}
+
+/**
+ * Send what the engine has waiting, once.
+ *
+ * Returns counts and, when the round stopped early, why (`stopped`: "busy", "no-engine",
+ * "engine-refused", "no-lease", "lease-refused", "unauthorized", "forbidden"), so the command
+ * can print it and the watcher can say it once. Throws only when this machine is not logged in,
+ * the way the journal upload does.
+ */
+async function sendPending({
+  limit = DEFAULT_LIMIT,
+  freshToken = api.ensureFreshToken,
+  fetchLease = lease.fetchLease,
+  fetchImpl = (...args) => fetch(...args),
+  log = () => {},
+} = {}) {
+  const result = {
+    lines: 0,
+    sent: 0,
+    documented: 0,
+    failed: 0,
+    waiting: 0,
+    stopped: null,
+    reason: null,
+    failures: [],
+    model: null,
+    imported: 0,
+    importFailed: null,
+    marked: null,
+  };
+  // Off in this release (src/lib/release.js `llmDocs`): nothing is asked of the engine, the service or the model host.
+  if (!require("./release").llmDocs) {
+    result.stopped = "off";
+    result.reason = NOT_IN_THIS_RELEASE;
+    return result;
+  }
+
+  const release = acquireLock();
+  if (!release) {
+    result.stopped = "busy";
+    return result;
+  }
+  try {
+    // 1. Is any folder built with --llm? Asked without a licence, so a machine that never used
+    //    the option makes no lease request, no login refresh and no emit, every ten minutes.
+    //    `null` means this engine cannot say (it predates `docs --marked`), and the round goes
+    //    on the old way: emit, which finds out for itself.
+    const marked = await markedFolders();
+    if (marked === "no-engine") {
+      result.stopped = "no-engine";
+      return result;
+    }
+    result.marked = marked;
+    if (marked === 0) return result;
+
+    // 2. Logged in, before emit, so a machine that cannot send never makes the engine read its
+    //    graph and build requests for nothing.
+    let token;
+    try {
+      token = await freshToken();
+    } catch (err) {
+      if (err && err.code === "ENOAUTH") {
+        throw new Error(`${err.message} Documentation waits on this machine until then.`);
+      }
+      throw err;
+    }
+
+    // 3. The engine's requests. Batches still waiting out a backoff are asked for on top of the
+    //    limit and then skipped, so a failing batch at the head of the engine's (sorted) list
+    //    never takes up the room the ones behind it need.
+    const now = Date.now();
+    const backoff = pruneBackoff(loadBackoff(), now);
+    const waitingCount = Object.keys(backoff).filter((b) => waitingUntil(backoff, b, now)).length;
+    const emitLimit = limit > 0 ? limit + waitingCount : 0;
+    const args = ["docs", "--emit", "--all"];
+    if (emitLimit > 0) args.push("--limit", String(emitLimit));
+    let emitted;
+    let emitNotes;
+    release.touch();
+    try {
+      ({ stdout: emitted, stderr: emitNotes } = await localEngine.runLicensed(args, {
+        log,
+        withStderr: true,
+      }));
+    } catch (err) {
+      const message = (err && err.message) || String(err);
+      if (/^cannot run /.test(message)) {
+        result.stopped = "no-engine";
+        return result;
+      }
+      // The engine's own words, which never carry a request: it prints counts and messages on
+      // standard error, and requests only on standard output (niro-local reuse/command.py).
+      result.stopped = "engine-refused";
+      result.reason = String(err.stderr || message).split("\n")[0].slice(0, 300);
+      return result;
+    }
+
+    const lines = [];
+    for (const raw of emitted.split("\n")) {
+      if (!raw.trim()) continue;
+      try {
+        const line = JSON.parse(raw);
+        if (line && line.request && typeof line.request === "object" && line.batch) lines.push(line);
+        else result.failed += 1;
+      } catch {
+        result.failed += 1;
+      }
+    }
+    result.lines = lines.length;
+    let due = lines.filter((line) => !waitingUntil(backoff, line.batch, now));
+    result.waiting = lines.length - due.length;
+    if (limit > 0) due = due.slice(0, limit);
+
+    // 4. The model, whenever a folder is built with --llm, even with nothing to send: without
+    //    it the engine can neither embed the notes it has nor a find_reusable_code question.
+    //    Fetched before the POSTs so the import at the end of this round can make vectors. A
+    //    failure here costs search, not documentation, so it is reported and the round goes on.
+    if (lines.length > 0 || marked > 0 || (marked === null && foldersMarked(emitNotes))) {
+      try {
+        result.model = await ensureModel({ fetchImpl, beforeEach: release.touch });
+      } catch (err) {
+        result.model = { error: (err && err.message) || "model fetch failed" };
+      }
+    }
+    if (due.length === 0) return result;
+
+    // 5. One lease for the round, in memory only.
+    let held = await getLease(fetchLease);
+    if (held.refused || held.unreachable) {
+      result.stopped = held.refused ? "lease-refused" : "no-lease";
+      result.reason = held.refused ? held.reason : held.detail;
+      return result;
+    }
+
+    // 6. Each batch, in turn.
+    const answers = [];
+    let answeredTotal = 0;
+    for (const line of due) {
+      release.touch();
+      if (leaseNearExpiry(held, Date.now())) {
+        const renewed = await getLease(fetchLease);
+        if (renewed.refused || renewed.unreachable) {
+          result.stopped = renewed.refused ? "lease-refused" : "no-lease";
+          result.reason = renewed.refused ? renewed.reason : renewed.detail;
+          break;
+        }
+        held = renewed;
+      }
+      const outcome = await postOne(line, held, token, fetchImpl);
+      const id = shortId(line.batch);
+      if (outcome.status === 401) {
+        // The token was good a moment ago; another try this round would be refused the same way.
+        result.stopped = "unauthorized";
+        result.reason = "the service did not accept this machine's login";
+        break;
+      }
+      if (outcome.status === 403) {
+        // About the plan or the lease, not this batch: no backoff for it, and nothing more is
+        // sent until the next round brings a new lease.
+        result.stopped = "forbidden";
+        result.reason = reasonOf(outcome.body, 403);
+        break;
+      }
+      result.sent += 1;
+      // Only 200: the engine stores an answer only when its status is exactly 200
+      // (niro-local reuse/answers.py, `int(line.get("status") or 0) != 200`).
+      const failedStatus = !outcome.error && outcome.status !== 200;
+      if (outcome.error || outcome.unreadable || failedStatus) {
+        // The status first: a 502 from a proxy usually has an HTML body, and "HTTP 502" is the
+        // useful half of that.
+        const why = outcome.error || (failedStatus ? `HTTP ${outcome.status}` : "unreadable answer");
+        result.failed += 1;
+        result.failures.push(`${id}: ${why}`);
+        recordFailure(backoff, line.batch, Date.now());
+        answers.push({ ...envelope(line), error: why });
+        log(`    documentation batch ${id}: ${why}`);
+        continue;
+      }
+      const answered = answeredCount(outcome.body, line.entities);
+      if (answered === 0) {
+        // A 200 that describes none of the methods: asked again right away, it would most
+        // likely come back empty again, so it waits like a failure does.
+        result.failed += 1;
+        result.failures.push(`${id}: answered none of ${Object.keys(line.entities || {}).length}`);
+        recordFailure(backoff, line.batch, Date.now());
+      } else {
+        answeredTotal += answered;
+        delete backoff[line.batch];
+      }
+      answers.push({ ...envelope(line), status: outcome.status, body: outcome.body });
+    }
+
+    // 7. The answers back to the engine, which stores what it can use.
+    if (answers.length > 0) {
+      release.touch();
+      try {
+        await localEngine.runLicensed(["docs", "--import"], {
+          stdin: answers.map((a) => JSON.stringify(a)).join("\n") + "\n",
+          log,
+        });
+        result.imported = answers.length;
+        // Counted only once the engine has them: an answer that never reached it documented
+        // nothing.
+        result.documented = answeredTotal;
+      } catch (err) {
+        // Not thrown: the backoff below must still be written. What was not stored is asked for
+        // again by the next emit, so nothing is lost but the call. The engine's standard error
+        // carries counts and messages only, never an answer.
+        result.importFailed = String((err && (err.stderr || err.message)) || "import failed")
+          .split("\n")[0]
+          .slice(0, 300);
+      }
+    }
+    try {
+      saveBackoff(backoff);
+    } catch (err) {
+      log(color.yellow(`    Could not write ${backoffFile()}: ${err.code || "write failed"}`));
+    }
+    return result;
+  } finally {
+    release();
+  }
+}
+
+// ---------------------------------------------------------------------------------------------
+// The watcher's timer
+// ---------------------------------------------------------------------------------------------
+
+/** About every ten minutes: often enough that a new folder is documented within the hour. */
+const LOCAL_DOCS_INTERVAL_MS = 10 * 60 * 1000;
+/** After the daemon's startup burst, like the other machine-wide jobs. */
+const LOCAL_DOCS_STARTUP_DELAY_MS = 3 * 60 * 1000;
+
+/**
+ * Send documentation requests on a timer, from the watch daemon.
+ *
+ * Machine-wide, so the daemon starts it once, beside the audit and journal uploads, and gated by
+ * its own periodicClaim name for the same two reasons: a PM2 restart must not turn a ten-minute
+ * job into one that runs on every restart, and a `niro watch` left in a terminal must not run it
+ * at the same moment (the run lock is the second guard for that).
+ *
+ * It stops itself on a machine with no engine and says nothing, like startLocalUploads. Anything
+ * else that stops a round is said once, and said again only when it changes, so a plan without
+ * documentation does not write the same sentence into the log every ten minutes.
+ *
+ * Never throws: the daemon's real job is keeping folders in sync.
+ */
+function startLocalDocs({
+  intervalMs = LOCAL_DOCS_INTERVAL_MS,
+  startupDelayMs = LOCAL_DOCS_STARTUP_DELAY_MS,
+  claimName = "localEngineDocs",
+  log = console.log,
+  send = () => module.exports.sendPending(),
+} = {}) {
+  let stopped = false;
+  let lastSaid = null;
+  const handle = { stop: () => {} };
+
+  const sayOnce = (text) => {
+    if (text === lastSaid) return;
+    lastSaid = text;
+    log(color.yellow(`    ${text}`));
+  };
+
+  const run = async () => {
+    if (stopped || !periodicClaim.claim(claimName, intervalMs)) return;
+    try {
+      const result = await send();
+      if (result.stopped === "no-engine") {
+        handle.stop();
+        return;
+      }
+      if (result.stopped === "busy") {
+        // Said every time it happens: a run that takes longer than the interval is worth seeing,
+        // and a lock that is never released shows up here rather than as silence.
+        log("    Documentation round skipped: another run holds the lock.");
+        return;
+      }
+      if (result.stopped) {
+        sayOnce(`Documentation not sent (${result.stopped}): ${result.reason || "no reason given"}`);
+        return;
+      }
+      if (result.documented > 0) {
+        lastSaid = null;
+        log(`    Documented ${result.documented} method(s) through Niro.`);
+      }
+      if (result.failed > 0) {
+        log(color.yellow(`    ${result.failed} documentation batch(es) failed and will be retried later.`));
+      }
+    } catch (err) {
+      sayOnce(`Could not send documentation requests: ${(err && err.message) || String(err)}`);
+    }
+  };
+
+  const first = setTimeout(run, startupDelayMs);
+  const timer = setInterval(run, intervalMs);
+  if (first.unref) first.unref();
+  if (timer.unref) timer.unref();
+  handle.stop = () => {
+    stopped = true;
+    clearTimeout(first);
+    clearInterval(timer);
+  };
+  return handle;
+}
+
+module.exports = {
+  NOT_IN_THIS_RELEASE,
+  sendPending,
+  ensureModel,
+  startLocalDocs,
+  acquireLock,
+  backoffFile,
+  lockFile,
+  ENRICH_PATH,
+  DEFAULT_LIMIT,
+  BACKOFF_BASE_MS,
+};