@runtypelabs/persona 3.21.3 → 3.23.0

package/src/utils/throughput-tracker.test.ts

@@ -0,0 +1,366 @@
+import { describe, it, expect } from "vitest";
+import {
+  ThroughputTracker,
+  estimateOutputTokens,
+  type ThroughputMetric,
+} from "./throughput-tracker";
+
+/**
+ * Builds a tracker with a manually-advanced clock so durations are
+ * deterministic. Call `clock.set(ms)` before feeding an event.
+ */
+function makeTracker() {
+  let nowMs = 0;
+  const tracker = new ThroughputTracker(() => nowMs);
+  return {
+    tracker,
+    at(ms: number) {
+      nowMs = ms;
+      return tracker;
+    },
+    metric(): ThroughputMetric {
+      return tracker.getMetric();
+    },
+  };
+}
+
+const text = (n: number) => "a".repeat(n);
+
+describe("estimateOutputTokens", () => {
+  it("uses the ~4 chars/token heuristic with a floor of 1", () => {
+    expect(estimateOutputTokens("")).toBe(0);
+    expect(estimateOutputTokens("   ")).toBe(0);
+    expect(estimateOutputTokens("a")).toBe(1);
+    expect(estimateOutputTokens(text(40))).toBe(10);
+  });
+});
+
+describe("ThroughputTracker — live estimate", () => {
+  it("estimates output tokens live from visible text deltas", () => {
+    const h = makeTracker();
+
+    // First delta: lazily starts the run. Duration is 0 so no rate yet.
+    h.at(1000).processEvent("step_delta", {
+      type: "step_delta",
+      text: text(40),
+    });
+    let m = h.metric();
+    expect(m.status).toBe("running");
+    expect(m.outputTokens).toBe(10);
+    expect(m.source).toBe("estimate");
+    expect(m.tokensPerSecond).toBeUndefined();
+
+    // Second delta 1s later via a different visible event type.
+    h.at(2000).processEvent("chunk", { type: "chunk", text: text(40) });
+    m = h.metric();
+    expect(m.outputTokens).toBe(20);
+    expect(m.durationMs).toBe(1000);
+    expect(m.source).toBe("estimate");
+    expect(m.tokensPerSecond).toBeCloseTo(20);
+  });
+
+  it("counts agent_turn_delta text deltas as visible output", () => {
+    const h = makeTracker();
+    h.at(0).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "text",
+      text: text(40),
+    });
+    h.at(1000).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "text",
+      text: text(40),
+    });
+    const m = h.metric();
+    expect(m.status).toBe("running");
+    expect(m.outputTokens).toBe(20);
+    expect(m.tokensPerSecond).toBeCloseTo(20);
+  });
+
+  it("ignores agent_turn_delta with a missing contentType (matches client)", () => {
+    const h = makeTracker();
+    // The client only renders agent text when contentType === "text"; a delta
+    // without a contentType is not visible output, so it must not be counted.
+    h.at(0).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      text: text(400),
+    });
+    expect(h.metric().status).toBe("idle");
+
+    // An explicit text delta is still counted.
+    h.at(1000).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "text",
+      text: text(40),
+    });
+    expect(h.metric().outputTokens).toBe(10);
+  });
+});
+
+describe("ThroughputTracker — exact usage finalization", () => {
+  it("prefers exact output tokens from the terminal event over the estimate", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    h.at(2000).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "text",
+      text: text(40),
+    });
+
+    // estimate so far would be 20 tokens; terminal usage overrides it.
+    h.at(2000).processEvent("flow_complete", {
+      type: "flow_complete",
+      usage: { outputTokens: 123 },
+    });
+
+    const m = h.metric();
+    expect(m.status).toBe("complete");
+    expect(m.outputTokens).toBe(123);
+    expect(m.source).toBe("usage");
+    expect(m.durationMs).toBe(1000); // streamed window 1000ms
+    expect(m.tokensPerSecond).toBeCloseTo(123);
+  });
+
+  it("accumulates exact usage from intermediate completes and uses it on terminal", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    h.at(1500).processEvent("step_complete", {
+      type: "step_complete",
+      result: { tokens: { output: 50 } },
+    });
+    h.at(2000).processEvent("agent_turn_complete", {
+      type: "agent_turn_complete",
+      usage: { outputTokens: 30 },
+    });
+
+    // Terminal carries no usage of its own → fall back to accumulated 80.
+    h.at(2000).processEvent("agent_complete", { type: "agent_complete" });
+
+    const m = h.metric();
+    expect(m.status).toBe("complete");
+    expect(m.outputTokens).toBe(80);
+    expect(m.source).toBe("usage");
+  });
+
+  it("falls back to provider execution time when the streamed window is too short", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("flow_start", { type: "flow_start" });
+    h.at(1100).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    // streamed window = 50ms (< 250ms threshold) → use executionTime.
+    h.at(1150).processEvent("flow_complete", {
+      type: "flow_complete",
+      executionTime: 5000,
+    });
+
+    const m = h.metric();
+    expect(m.status).toBe("complete");
+    expect(m.outputTokens).toBe(10);
+    expect(m.durationMs).toBe(5000);
+    expect(m.tokensPerSecond).toBeCloseTo(2);
+  });
+});
+
+describe("ThroughputTracker — non-visible deltas ignored", () => {
+  it("ignores agent thinking and tool_input deltas", () => {
+    const h = makeTracker();
+
+    h.at(1000).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "thinking",
+      text: text(400),
+    });
+    expect(h.metric().status).toBe("idle");
+
+    h.at(1000).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "tool_input",
+      text: text(400),
+    });
+    expect(h.metric().status).toBe("idle");
+
+    // Only the visible text delta is counted.
+    h.at(1000).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "text",
+      text: text(40),
+    });
+    h.at(2000).processEvent("agent_turn_delta", {
+      type: "agent_turn_delta",
+      contentType: "thinking",
+      text: text(400),
+    });
+
+    const m = h.metric();
+    expect(m.status).toBe("running");
+    expect(m.outputTokens).toBe(10); // thinking text excluded
+  });
+
+  it("ignores tool and context step deltas", () => {
+    const h = makeTracker();
+
+    h.at(1000).processEvent("step_delta", {
+      type: "step_delta",
+      stepType: "tool",
+      text: text(400),
+    });
+    expect(h.metric().status).toBe("idle");
+
+    h.at(1000).processEvent("step_delta", {
+      type: "step_delta",
+      executionType: "context",
+      text: text(400),
+    });
+    expect(h.metric().status).toBe("idle");
+
+    // A prompt-step delta is counted.
+    h.at(1000).processEvent("step_delta", {
+      type: "step_delta",
+      stepType: "prompt",
+      text: text(40),
+    });
+    expect(h.metric().status).toBe("running");
+    expect(h.metric().outputTokens).toBe(10);
+  });
+});
+
+describe("ThroughputTracker — intermediate completes do not finalize", () => {
+  it("keeps the run running across step_complete / agent_turn_complete", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+
+    h.at(1500).processEvent("step_complete", {
+      type: "step_complete",
+      result: { tokens: { output: 10 } },
+    });
+    expect(h.metric().status).toBe("running");
+
+    h.at(1800).processEvent("agent_turn_complete", {
+      type: "agent_turn_complete",
+      usage: { outputTokens: 5 },
+    });
+    expect(h.metric().status).toBe("running");
+
+    h.at(2000).processEvent("flow_complete", { type: "flow_complete" });
+    expect(h.metric().status).toBe("complete");
+  });
+});
+
+describe("ThroughputTracker — error handling", () => {
+  it.each(["step_error", "flow_error", "agent_error", "error"])(
+    "marks the metric unavailable on %s",
+    (errorType: string) => {
+      const h = makeTracker();
+      h.at(1000).processEvent("step_delta", {
+        type: "step_delta",
+        text: text(40),
+      });
+      expect(h.metric().status).toBe("running");
+
+      h.at(1500).processEvent(errorType, { type: errorType });
+      expect(h.metric().status).toBe("error");
+      expect(h.metric().tokensPerSecond).toBeUndefined();
+    }
+  );
+
+  it("treats a bare non-object error payload as an error", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    h.at(1500).processEvent("error", "boom");
+    expect(h.metric().status).toBe("error");
+  });
+
+  it("ignores terminal/error events when no run is active", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("flow_complete", { type: "flow_complete" });
+    expect(h.metric().status).toBe("idle");
+    h.at(1000).processEvent("flow_error", { type: "flow_error" });
+    expect(h.metric().status).toBe("idle");
+  });
+});
+
+describe("ThroughputTracker — reset & re-run", () => {
+  it("reset() returns to idle", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    h.tracker.reset();
+    expect(h.metric()).toEqual({ status: "idle" });
+  });
+
+  it("starts a fresh run after a completed one", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(80) });
+    h.at(2000).processEvent("flow_complete", { type: "flow_complete" });
+    expect(h.metric().status).toBe("complete");
+
+    // New stream → accumulation resets, does not carry the prior 20 tokens.
+    h.at(3000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    const m = h.metric();
+    expect(m.status).toBe("running");
+    expect(m.outputTokens).toBe(10);
+  });
+
+  it("resets a stale run on the next request's start event (no bleed)", () => {
+    const h = makeTracker();
+    // First request streams visible output but never terminates (e.g. the
+    // user cancels mid-stream — no flow_complete / error frame is emitted).
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(120) });
+    expect(h.metric().outputTokens).toBe(30);
+    expect(h.metric().status).toBe("running");
+
+    // Next request begins. Its flow_start must discard the stale run so the
+    // prior request's 30 tokens don't bleed into this one.
+    h.at(5000).processEvent("flow_start", { type: "flow_start" });
+    h.at(5200).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    const m = h.metric();
+    expect(m.status).toBe("running");
+    expect(m.outputTokens).toBe(10); // only the new request's text
+  });
+
+  it("does not reset between per-step starts within one request", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("flow_start", { type: "flow_start" });
+    h.at(1100).processEvent("step_start", { type: "step_start" });
+    h.at(1200).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    // A second step within the same request must not restart accumulation.
+    h.at(1300).processEvent("step_start", { type: "step_start" });
+    h.at(1400).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    expect(h.metric().outputTokens).toBe(20);
+  });
+});
+
+describe("ThroughputTracker — exact usage never drops mid-run", () => {
+  it("keeps exact tokens as a floor when later steps stream more text", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+
+    // Step 1 reports exact usage; the running total switches to it.
+    h.at(1500).processEvent("step_complete", {
+      type: "step_complete",
+      result: { tokens: { output: 50 } },
+    });
+    let m = h.metric();
+    expect(m.outputTokens).toBe(50);
+    expect(m.source).toBe("usage");
+
+    // Step 2 streams more visible text — the total must grow from 50, not
+    // collapse back to a bare 10-token estimate of the new text.
+    h.at(2000).processEvent("step_delta", { type: "step_delta", text: text(40) });
+    m = h.metric();
+    expect(m.outputTokens).toBe(60); // 50 exact + 10 estimated
+    expect(m.source).toBe("usage");
+  });
+});
+
+describe("ThroughputTracker — live rate decays while paused", () => {
+  it("recomputes duration/tok-s from the clock between events", () => {
+    const h = makeTracker();
+    h.at(1000).processEvent("step_delta", { type: "step_delta", text: text(400) });
+    // 100 tokens over a 1s window read at t=2000 → ~100 tok/s.
+    h.at(2000);
+    expect(h.metric().tokensPerSecond).toBeCloseTo(100);
+    // Same tokens, but the model has paused — reading at t=5000 (4s window)
+    // must decay the displayed rate even though no new event arrived.
+    h.at(5000);
+    expect(h.metric().tokensPerSecond).toBeCloseTo(25);
+  });
+});