test/antigravity-usage-baseline.test.js

Rama main · réplica de lectura del repositorio original.

var test = require("node:test");
var assert = require("node:assert/strict");
var EventEmitter = require("events");
var PassThrough = require("stream").PassThrough;
var antigravity = require("../lib/yoke/adapters/antigravity");

// turns[i] = { status, usage, duration, steps: [{ usage, duration }] }
function scriptedSpawn(turns) {
  var turn = 0;
  return function () {
    var proc = new EventEmitter();
    proc.stdout = new PassThrough();
    proc.stderr = new PassThrough();
    proc.stdin = new EventEmitter();
    proc.stdin.write = function () {
      var spec = turns[turn++];
      setImmediate(function () {
        (spec.steps || []).forEach(function (step, index) {
          proc.stdout.write(JSON.stringify({ event: "step_update", step_update: {
            step_index: index, state: "DONE", step_type: "agent_response", text_delta: "t" + turn,
            usage: step.usage, duration_seconds: step.duration,
          } }) + "\n");
        });
        var result = { status: spec.status || "SUCCESS", conversation_id: "conv" };
        if (spec.usage) result.usage = spec.usage;
        if (typeof spec.duration === "number") result.duration_seconds = spec.duration;
        proc.stdout.write(JSON.stringify({ event: "result", result: result }) + "\n");
      });
      return true;
    };
    proc.stdin.end = function () { setImmediate(function () { proc.emit("exit", 0, null); }); };
    proc.kill = function () { proc.emit("exit", null, "SIGINT"); };
    return proc;
  };
}

function usage(input, output, cache) {
  return { input_tokens: input, output_tokens: output, cache_read_tokens: cache || 0 };
}

async function runTurns(turns, opts) {
  var handle = antigravity.createAntigravityQueryHandle("/contract/agy", Object.assign({ cwd: "/project", _spawn: scriptedSpawn(turns) }, opts || {}), function () {});
  var iterator = handle[Symbol.asyncIterator]();
  var results = [];
  for (var i = 0; i < turns.length; i++) {
    handle.pushMessage("turn " + (i + 1));
    if (i === turns.length - 1) handle.endInput();
    for (;;) {
      var next = await iterator.next();
      if (next.done) break;
      if (next.value.yokeType === "result") {
        results.push({
          usage: next.value.usage ? [next.value.usage.input_tokens, next.value.usage.output_tokens, next.value.usage.cache_read_input_tokens] : null,
          duration: next.value.duration,
        });
        break;
      }
    }
  }
  return results;
}

test("a turn without counters falls back to step totals and the next result re-baselines", async function () {
  var results = await runTurns([
    { status: "CANCELED", steps: [{ usage: usage(10, 2), duration: 1 }] },
    { usage: usage(14, 5), duration: 3, steps: [{ usage: usage(4, 3), duration: 0.5 }] },
    { usage: usage(20, 9), duration: 4.5 },
  ]);
  assert.deepEqual(results, [
    { usage: [10, 2, 0], duration: 1000 },
    { usage: [4, 3, 0], duration: 500 },
    { usage: [6, 4, 0], duration: 1500 },
  ]);
});

test("intermediate missing counters report unknown instead of historical totals", async function () {
  var results = await runTurns([
    { usage: usage(10, 2, 1), duration: 2 },
    { status: "ERROR" },
    { usage: usage(30, 10, 4), duration: 9 },
    { usage: usage(33, 12, 4), duration: 10 },
  ]);
  assert.deepEqual(results, [
    { usage: [10, 2, 1], duration: 2000 },
    { usage: null, duration: null },
    { usage: null, duration: null },
    { usage: [3, 2, 0], duration: 1000 },
  ]);
});

test("error and cancelled turns with counters keep a trusted baseline", async function () {
  var results = await runTurns([
    { usage: usage(10, 2), duration: 1 },
    { status: "ERROR", usage: usage(12, 3), duration: 1.5 },
    { status: "CANCELED", usage: usage(15, 5), duration: 2 },
    { usage: usage(16, 7), duration: 2.25 },
  ]);
  assert.deepEqual(results.map(function (r) { return r.usage; }), [[10, 2, 0], [2, 1, 0], [3, 2, 0], [1, 2, 0]]);
  assert.deepEqual(results.map(function (r) { return r.duration; }), [1000, 500, 500, 250]);
});

test("counters that go backwards are not trusted as a baseline", async function () {
  var results = await runTurns([
    { usage: usage(50, 20), duration: 5 },
    { usage: usage(8, 3), duration: 1, steps: [{ usage: usage(8, 3), duration: 1 }] },
    { usage: usage(9, 3), duration: 0.5 },
    { usage: usage(12, 6), duration: 2 },
  ]);
  assert.deepEqual(results, [
    { usage: [50, 20, 0], duration: 5000 },
    { usage: [8, 3, 0], duration: 1000 },
    { usage: [1, 0, 0], duration: null },
    { usage: [3, 3, 0], duration: 1500 },
  ]);
});

test("a resumed conversation starts unknown, uses step totals when present, and recovers", async function () {
  var unknown = await runTurns([
    { usage: usage(100, 50, 20), duration: 10 },
    { status: "CANCELED" },
    { usage: usage(110, 55, 21), duration: 12 },
    { usage: usage(111, 57, 21), duration: 12.5 },
  ], { resumeSessionId: "conv" });
  assert.deepEqual(unknown, [
    { usage: null, duration: null },
    { usage: null, duration: null },
    { usage: null, duration: null },
    { usage: [1, 2, 0], duration: 500 },
  ]);
  var measured = await runTurns([
    { usage: usage(100, 50), duration: 10, steps: [{ usage: usage(3, 1), duration: 0.25 }, { usage: usage(2, 1), duration: 0.5 }] },
    { usage: usage(104, 53), duration: 11 },
  ], { resumeSessionId: "conv" });
  assert.deepEqual(measured, [
    { usage: [5, 2, 0], duration: 750 },
    { usage: [4, 3, 0], duration: 1000 },
  ]);
});