@hicaru/pi-rlm 0.1.0 → 0.1.2

package/src/telemetry/mlflow-sink.ts

@@ -1,136 +0,0 @@
-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

@@ -1,99 +0,0 @@
-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

@@ -1,8 +0,0 @@
-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/tool/apply-diff-tool.ts

@@ -1,125 +0,0 @@
-/**
- * `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);
-    },
-  };
-}