@hicaru/pi-rlm 0.1.0

package/src/telemetry/mlflow-sink.ts

@@ -0,0 +1,136 @@
+import type { SubcallKind } from "../tool/rlm-details.ts";
+import type { SubcallStart } from "../state/events.ts";
+import { Dispatcher, type DispatcherSink } from "./dispatcher.ts";
+import {
+  MlflowTracer,
+  msToNs,
+  SpanStatusCode,
+  SpanType,
+  type MlflowConfig,
+  type SpanTracer,
+} from "./mlflow.ts";
+import type { TelemetrySink } from "./sink.ts";
+
+// @mlflow/core keeps process-global trace state. RLM telemetry uses one shared
+// tracer per process and should not be combined with another MLflow init path in
+// the same host process.
+let sharedTracer: MlflowTracer | undefined;
+function tracerFor(config: MlflowConfig): MlflowTracer {
+  sharedTracer ??= new MlflowTracer(config);
+  return sharedTracer;
+}
+
+const SPAN_TYPE: Readonly<Record<SubcallKind, SpanType>> = Object.freeze({
+  root: SpanType.AGENT,
+  rlm: SpanType.AGENT,
+  llm: SpanType.CHAT_MODEL,
+  batch: SpanType.CHAT_MODEL,
+  tool: SpanType.TOOL,
+});
+
+type SpanOp =
+  | { readonly kind: "start"; readonly id: string; readonly info: SubcallStart; readonly tsMs: number }
+  | { readonly kind: "usage"; readonly id: string; readonly costUsd: number; readonly tokens: number }
+  | { readonly kind: "end"; readonly id: string; readonly error?: string; readonly resultPreview?: string; readonly tsMs: number };
+
+class SpanApplier<S> implements DispatcherSink<SpanOp> {
+  readonly name = "rlm-mlflow";
+  private readonly spans = new Map<string, S>();
+  private readonly usage = new Map<string, { cost: number; tokens: number }>();
+
+  constructor(private readonly tracer: SpanTracer<S>) {}
+
+  async handle(op: SpanOp): Promise<void> {
+    if (op.kind === "start") {
+      const parent = op.info.parentId ? this.spans.get(op.info.parentId) : undefined;
+      const span = this.tracer.startSpan({
+        name: op.info.label,
+        spanType: SPAN_TYPE[op.info.kind],
+        parent,
+        inputs: {
+          ...(op.info.model ? { model: op.info.model } : {}),
+          ...(op.info.detail ? { detail: op.info.detail } : {}),
+        },
+        attributes: {
+          "rlm.kind": op.info.kind,
+          "rlm.depth": op.info.depth,
+          ...(op.info.model ? { "rlm.model": op.info.model } : {}),
+          ...(op.info.runId ? { "rlm.runId": op.info.runId } : {}),
+          ...(op.info.resume ? { "rlm.resume": true } : {}),
+        },
+        startTimeNs: msToNs(op.tsMs),
+      });
+      if (span !== undefined) this.spans.set(op.id, span);
+      return;
+    }
+
+    if (op.kind === "usage") {
+      const usage = this.usage.get(op.id) ?? { cost: 0, tokens: 0 };
+      this.usage.set(op.id, { cost: usage.cost + op.costUsd, tokens: usage.tokens + op.tokens });
+      return;
+    }
+
+    this.closeSpan(op.id, {
+      outputs: op.resultPreview ? { result: op.resultPreview } : undefined,
+      status: op.error ? SpanStatusCode.ERROR : undefined,
+      endTimeNs: msToNs(op.tsMs),
+    });
+  }
+
+  async flush(): Promise<void> {
+    await this.tracer.flush();
+  }
+
+  async shutdown(): Promise<void> {
+    for (const id of [...this.spans.keys()]) this.closeSpan(id, { status: SpanStatusCode.ERROR });
+    await this.tracer.flush();
+  }
+
+  private closeSpan(
+    id: string,
+    opts: { readonly outputs?: Readonly<Record<string, unknown>>; readonly status?: SpanStatusCode; readonly endTimeNs?: number },
+  ): void {
+    const span = this.spans.get(id);
+    if (span === undefined) return;
+    const usage = this.usage.get(id);
+    this.tracer.endSpan(span, {
+      ...(usage ? { attributes: { "rlm.usage.cost_usd": usage.cost, "rlm.usage.tokens": usage.tokens } } : {}),
+      ...opts,
+    });
+    this.spans.delete(id);
+    this.usage.delete(id);
+  }
+}
+
+export class MlflowSink implements TelemetrySink {
+  private readonly dispatcher: Dispatcher<SpanOp>;
+
+  constructor(config: MlflowConfig, maxQueueSize = 100, tracer: SpanTracer = tracerFor(config)) {
+    this.dispatcher = new Dispatcher<SpanOp>({ maxQueueSize });
+    this.dispatcher.registerSink(new SpanApplier(tracer));
+  }
+
+  start(id: string, info: SubcallStart): void {
+    if (id) this.dispatcher.dispatch({ kind: "start", id, info, tsMs: Date.now() });
+  }
+
+  usage(id: string, costUsd: number, tokens: number): void {
+    if (id) this.dispatcher.dispatch({ kind: "usage", id, costUsd, tokens });
+  }
+
+  end(id: string, opts?: { readonly error?: string; readonly resultPreview?: string }): void {
+    if (id) {
+      this.dispatcher.dispatch({
+        kind: "end",
+        id,
+        error: opts?.error,
+        resultPreview: opts?.resultPreview,
+        tsMs: Date.now(),
+      });
+    }
+  }
+
+  shutdown(): Promise<void> {
+    return this.dispatcher.shutdown();
+  }
+}

package/src/telemetry/mlflow.ts

@@ -0,0 +1,99 @@
+import { flushTraces, init, type LiveSpan, startSpan as mlflowStartSpan, SpanStatusCode, SpanType } from "@mlflow/core";
+import { resolveMlflowConfig, type MlflowConfig } from "./mlflow-config.ts";
+
+export { type LiveSpan, type MlflowConfig, SpanStatusCode, SpanType };
+
+export function msToNs(ms: number): number {
+  return ms * 1_000_000;
+}
+
+export interface StartSpanOptions<S> {
+  readonly name: string;
+  readonly spanType: SpanType;
+  readonly parent?: S;
+  readonly inputs?: Readonly<Record<string, unknown>>;
+  readonly attributes?: Readonly<Record<string, unknown>>;
+  readonly startTimeNs?: number;
+}
+
+export interface EndSpanOptions {
+  readonly attributes?: Readonly<Record<string, unknown>>;
+  readonly outputs?: Readonly<Record<string, unknown>>;
+  readonly status?: SpanStatusCode;
+  readonly endTimeNs?: number;
+}
+
+export interface SpanTracer<S = unknown> {
+  startSpan(options: StartSpanOptions<S>): S | undefined;
+  endSpan(span: S, options?: EndSpanOptions): void;
+  flush(): Promise<void>;
+  shutdown(): Promise<void>;
+}
+
+export class MlflowTracer implements SpanTracer<LiveSpan> {
+  private initialized = false;
+  private initAttempted = false;
+
+  constructor(private readonly config: MlflowConfig) {}
+
+  startSpan(options: StartSpanOptions<LiveSpan>): LiveSpan | undefined {
+    if (!this.ensureInit()) return undefined;
+    const span = mlflowStartSpan({
+      name: options.name,
+      spanType: options.spanType,
+      ...(options.parent ? { parent: options.parent } : {}),
+      ...(options.inputs ? { inputs: options.inputs } : {}),
+      ...(options.startTimeNs !== undefined ? { startTimeNs: options.startTimeNs } : {}),
+    });
+    if (options.attributes) {
+      for (const [key, value] of Object.entries(options.attributes)) span.setAttribute(key, value);
+    }
+    return span;
+  }
+
+  endSpan(span: LiveSpan, options: EndSpanOptions = {}): void {
+    if (options.attributes) {
+      for (const [key, value] of Object.entries(options.attributes)) span.setAttribute(key, value);
+    }
+    span.end({
+      ...(options.outputs ? { outputs: options.outputs } : {}),
+      ...(options.status !== undefined ? { status: options.status } : {}),
+      ...(options.endTimeNs !== undefined ? { endTimeNs: options.endTimeNs } : {}),
+    });
+  }
+
+  async flush(): Promise<void> {
+    if (!this.initialized) return;
+    try {
+      await flushTraces();
+    } catch (err) {
+      console.warn(`[rlm-telemetry] mlflow flush failed: ${err instanceof Error ? err.message : String(err)}`);
+    }
+  }
+
+  async shutdown(): Promise<void> {
+    await this.flush();
+  }
+
+  private ensureInit(): boolean {
+    if (this.initialized || this.initAttempted) return this.initialized;
+    this.initAttempted = true;
+    const resolved = resolveMlflowConfig(this.config);
+    if (!resolved.trackingUri) {
+      console.warn("[rlm-telemetry] MLFLOW_TRACKING_URI unset — spans dropped");
+      return false;
+    }
+
+    try {
+      init({
+        trackingUri: resolved.trackingUri,
+        experimentId: resolved.experimentId ?? "0",
+        ...(resolved.trackingToken ? { trackingServerToken: resolved.trackingToken } : {}),
+      });
+      this.initialized = true;
+    } catch (err) {
+      console.warn(`[rlm-telemetry] mlflow init failed: ${err instanceof Error ? err.message : String(err)}`);
+    }
+    return this.initialized;
+  }
+}

package/src/telemetry/sink.ts

@@ -0,0 +1,8 @@
+import type { SubcallStart } from "../state/events.ts";
+
+export interface TelemetrySink {
+  start(id: string, info: SubcallStart): void;
+  usage(id: string, costUsd: number, tokens: number): void;
+  end(id: string, opts?: { readonly error?: string; readonly resultPreview?: string }): void;
+  shutdown(): Promise<void>;
+}

package/src/text/edits.ts

@@ -0,0 +1,16 @@
+export interface AnchorEdit {
+  readonly oldText: string;
+  readonly newText: string;
+}
+
+export function countOccurrences(haystack: string, needle: string): number {
+  if (needle.length === 0) return 0;
+  let count = 0;
+  let offset = 0;
+  for (;;) {
+    const match = haystack.indexOf(needle, offset);
+    if (match < 0) return count;
+    count++;
+    offset = match + needle.length;
+  }
+}

package/src/text/parsing.ts

@@ -0,0 +1,35 @@
+/**
+ * Parsing helpers: extract ```repl``` code blocks from a model response.
+ *
+ * The RLM root model emits Python wrapped in fenced blocks tagged `repl`. We extract those
+ * blocks in order; everything else is prose the model uses to think out loud.
+ */
+
+const FENCE = /```[ \t]*repl[ \t]*\r?\n([\s\S]*?)```/g;
+
+/** Return every ```repl``` block body, in document order. */
+export function findReplBlocks(text: string): string[] {
+  const blocks: string[] = [];
+  let m: RegExpExecArray | null;
+  FENCE.lastIndex = 0;
+  while ((m = FENCE.exec(text)) !== null) {
+    const code = m[1] ?? "";
+    if (code.trim()) blocks.push(code.replace(/\s+$/, ""));
+  }
+  return blocks;
+}
+
+/** True if the response contains at least one runnable ```repl``` block. */
+export function hasReplBlock(text: string): boolean {
+  FENCE.lastIndex = 0;
+  return FENCE.test(text);
+}
+
+/** Truncate REPL stdout for the model's context window (head + tail, with an elision note). */
+export function truncateOutput(text: string, limit = 20_000): string {
+  if (text.length <= limit) return text;
+  const head = Math.floor(limit * 0.7);
+  const tail = limit - head;
+  const cut = text.length - head - tail;
+  return `${text.slice(0, head)}\n... [${cut} chars elided] ...\n${text.slice(-tail)}`;
+}

package/src/text/preview.ts

@@ -0,0 +1,18 @@
+/** Shared short text previews for live tree rows. */
+
+import type { ReplResult } from "../sandbox/protocol.ts";
+
+const DEFAULT_PREVIEW_CHARS = 200;
+
+export function previewText(text: string, maxChars = DEFAULT_PREVIEW_CHARS): string {
+  const normalized = text.replace(/\s+/g, " ").trim();
+  return normalized.length > maxChars ? `${normalized.slice(0, Math.max(0, maxChars - 1))}…` : normalized;
+}
+
+export function previewStdout(results: readonly ReplResult[]): string {
+  for (let index = results.length - 1; index >= 0; index--) {
+    const stdout = results[index]?.stdout.trim();
+    if (stdout) return previewText(stdout);
+  }
+  return "";
+}

package/src/text/tokens.ts

@@ -0,0 +1,64 @@
+/**
+ * Lightweight, dependency-free token estimation.
+ *
+ * We deliberately avoid a tokenizer dependency: RLM only needs rough budgets to decide when
+ * to chunk or compact, and a ~4-chars/token heuristic is accurate enough for that. Real token
+ * accounting comes back from the provider in `usage` after each call.
+ */
+
+const CHARS_PER_TOKEN = 4;
+
+/** Rough token count for a list of role/content messages. */
+export function estimateMessageTokens(messages: { content: string }[]): number {
+  let chars = 0;
+  for (const m of messages) chars += m.content.length + 8; // small per-message overhead
+  return Math.ceil(chars / CHARS_PER_TOKEN);
+}
+
+/** Total character length of a context payload (string or list of strings). */
+export function contextLength(context: unknown): number {
+  if (typeof context === "string") return context.length;
+  if (Array.isArray(context)) return context.reduce<number>((n, x) => n + String(x).length, 0);
+  return JSON.stringify(context ?? "").length;
+}
+
+/** Human label for a context payload's type, used in the metadata prompt. */
+export function contextTypeLabel(context: unknown): string {
+  if (typeof context === "string") return "str";
+  if (Array.isArray(context)) return `list[${context.length}]`;
+  return typeof context;
+}
+
+/** Compact per-file token distribution for a bundle context (the article's `context_lengths`). */
+export interface ContextSizeStats {
+  readonly files: number;
+  readonly min: number;
+  readonly median: number;
+  readonly max: number;
+}
+
+/** `true` if `v` is a context entry carrying an estimated `tokens` count. */
+const isTokenizedEntry = (v: unknown): v is { readonly tokens: number } =>
+  typeof v === "object" && v !== null && typeof (v as { readonly tokens?: unknown }).tokens === "number";
+
+/** Per-file token distribution for a context payload; `undefined` for plain strings or empty arrays.
+ *  Handles both serialized ContextFile[] (flat array from serializeForSandbox) and raw ContextBundle
+ *  objects ({ files: [...] }) so callers don't need to know which form they received. */
+export function contextSizeStats(context: unknown): ContextSizeStats | undefined {
+  // Normalise to a flat entry list: accept either a direct array or an object with a .files array.
+  const entries: readonly unknown[] = Array.isArray(context)
+    ? context
+    : Array.isArray((context as { readonly files?: unknown } | null)?.files)
+      ? (context as { readonly files: readonly unknown[] }).files
+      : [];
+  if (entries.length === 0) return undefined;
+  const sizes = new Array<number>(entries.length);
+  for (let i = 0; i < entries.length; i++) {
+    const entry = entries[i];
+    sizes[i] = isTokenizedEntry(entry) ? entry.tokens : 0;
+  }
+  sizes.sort((a, b) => a - b);
+  const mid = sizes.length >> 1;
+  const median = sizes.length % 2 !== 0 ? sizes[mid] : Math.round((sizes[mid - 1] + sizes[mid]) / 2);
+  return Object.freeze<ContextSizeStats>({ files: sizes.length, min: sizes[0], median, max: sizes[sizes.length - 1] });
+}

package/src/tool/apply-diff-tool.ts

@@ -0,0 +1,125 @@
+/**
+ * `apply_diff` — the main-agent-facing editing tool.
+ *
+ * Lightweight by design: the root model produces a complete unified diff
+ * directly and this tool validates its header then writes to disk via `applyEdits()`.
+ * Unlike the old `propose_edits`,
+ * there is NO inner RLM engine turn (no generate→validate→revise loop) — the
+ * model is trusted to emit a correct diff.
+ */
+
+import { Type } from "typebox";
+import { Text } from "@earendil-works/pi-tui";
+import type { ToolDefinition } from "@earendil-works/pi-coding-agent";
+import { applyEdits } from "../patch/index.ts";
+import { validateToolParams } from "./tool-utils.ts";
+import { errorMessage } from "../util/errors.ts";
+import { headlineStatusGlyph } from "./subcall-render.ts";
+import type { ProposedDiffEdit } from "../sandbox/protocol.ts";
+
+// ── Parameter schema ──────────────────────────────────────────────────────
+
+const ApplyDiffParams = Object.freeze(Type.Object({
+  diff: Type.String({
+    description: "Full unified diff string. Must include --- a/<path> / +++ b/<path> header and @@ hunk markers.",
+  }),
+}));
+
+// ── Details (no inner engine → no sub-call tree) ──────────────────────────
+
+export interface ApplyDiffDetails {
+  readonly status: "done" | "error";
+  readonly path: string;
+}
+
+const errorDetails = (): ApplyDiffDetails => ({ status: "error", path: "" });
+
+// ── Helpers ───────────────────────────────────────────────────────────────
+
+/** True when the diff carries a real `--- a/<path>` header (not a bare `---`). */
+function hasValidHeader(diff: string): boolean {
+  return /^--- \S/m.test(diff);
+}
+
+/** Extract the first target file path from a `--- a/<path>` header line. */
+function firstDiffPath(diff: string): string {
+  const match = /^--- a\/(.+)$/m.exec(diff);
+  return match?.[1] ?? "(unknown path)";
+}
+
+/** First ~60 chars of the diff body, starting at the first hunk marker. */
+function bodyPreview(diff: string): string {
+  const start = diff.indexOf("@@");
+  const body = start >= 0 ? diff.slice(start) : diff;
+  return body.replace(/\n/g, " ").slice(0, 60);
+}
+
+// ── Tool factory ──────────────────────────────────────────────────────────
+
+export function createApplyDiffTool(): ToolDefinition<typeof ApplyDiffParams, ApplyDiffDetails> {
+  return {
+    name: "apply_diff",
+    label: "Apply Diff",
+    description: [
+      "Apply a complete unified diff directly to disk.",
+      "The diff MUST include a --- a/<path> / +++ b/<path> header and @@ hunk markers.",
+         ].join(" "),
+    parameters: ApplyDiffParams,
+
+    async execute(_toolCallId, rawParams, _signal, _onUpdate, ctx) {
+      const validation = validateToolParams(ApplyDiffParams, rawParams, "apply_diff", errorDetails);
+      if (!validation.ok) return validation.error;
+      const { diff } = validation.value;
+
+      if (!diff.trim()) {
+        return {
+          content: [{ type: "text", text: "apply_diff requires a non-empty diff." }],
+          details: errorDetails(),
+        };
+      }
+
+      // No `path` param → the model MUST include a complete header. Reject with a
+      // clear message instead of a silent fallback so a malformed diff is surfaced.
+      if (!hasValidHeader(diff)) {
+        return {
+          content: [{
+            type: "text",
+            text: "apply_diff requires a complete unified diff with --- a/<path> / +++ b/<path> header and @@ hunk markers.",
+          }],
+          details: errorDetails(),
+        };
+      }
+
+      const path = firstDiffPath(diff);
+      const diffEdit: ProposedDiffEdit = { diff };
+      try {
+        await applyEdits([], [diffEdit], ctx);
+      } catch (e) {
+        return {
+          content: [{ type: "text", text: `apply_diff failed: ${errorMessage(e)}` }],
+          details: { status: "error", path },
+        };
+      }
+
+      return {
+        content: [{ type: "text", text: `Applied diff to ${path}.` }],
+        details: { status: "done", path },
+      };
+    },
+
+    renderCall(args, theme) {
+      const path = firstDiffPath(args.diff);
+      return new Text(
+        [theme.fg("toolTitle", theme.bold("apply_diff ")), theme.fg("dim", `${path}: ${bodyPreview(args.diff)}`)].join(""),
+        0, 0,
+      );
+    },
+
+    renderResult(result, _options, theme) {
+      const details = result.details as ApplyDiffDetails | undefined;
+      const path = details?.path ? details.path : "(unknown path)";
+      const glyph = headlineStatusGlyph(details?.status === "error" ? "error" : "done", theme);
+      return new Text(`${glyph} ${theme.fg("toolTitle", theme.bold("apply_diff"))} · ${path}`, 0, 0);
+    },
+  };
+}

package/src/tool/emitter-listener.ts

@@ -0,0 +1,24 @@
+/** Shared unsubscribe lifecycle for classes that attach to RlmEmitter events. */
+
+import { errorMessage } from "../util/errors.ts";
+
+export abstract class EmitterListener {
+  private readonly unsubs: (() => void)[] = [];
+
+  protected track(unsub: () => void): void {
+    this.unsubs.push(unsub);
+  }
+
+  protected trackAll(unsubs: readonly (() => void)[]): void {
+    for (const unsub of unsubs) this.track(unsub);
+  }
+
+  /** Detach all registered listeners. Safe to call more than once. */
+  dispose(): void {
+    const unsubs = this.unsubs.splice(0);
+    for (const unsub of unsubs) {
+      try { unsub(); }
+      catch (error) { console.warn(`[rlm] listener cleanup failed: ${errorMessage(error)}`); }
+    }
+  }
+}

package/src/tool/repl-details.ts

@@ -0,0 +1,23 @@
+/**
+ * ReplDetails — structured payload for the repl() tool's AgentToolResult<T>.
+ *
+ * Mirrors RlmDetails but scoped to a single code execution. Sub-calls (llm_query,
+ * rlm_query, todo, ask_user_question) triggered during sandbox execution are
+ * accumulated into the subcalls array for tree rendering.
+ */
+
+import type { RlmSubcall } from "./rlm-details.ts";
+
+export interface ReplDetails {
+  readonly status: "running" | "done" | "error";
+  /** stdout from the Python execution. */
+  readonly output: string;
+  /** stderr from the Python execution. */
+  readonly stderr: string;
+  /** Wall-clock execution time in milliseconds. */
+  readonly executionTimeMs: number;
+  /** Sub-calls triggered during this execution (llm_query, rlm_query, todo, etc.). */
+  readonly subcalls: readonly RlmSubcall[];
+  /** Running totals for this repl() call (cost + tokens from sub-LLM calls). */
+  readonly totals: { readonly costUsd: number; readonly tokens: number };
+}