Repository Public
1/** Governors: turn / wall / money limits with the final summarising turn, context guard, stall detection, cancel, steering. */2import { postOwnerMessage, requestCancel, resolveWait } from '@aos/kernel';3import { afterAll, afterEach, beforeAll, describe, expect, it } from 'vitest';4import { budgetView, finalTurnDue, StallDetector } from '../src/index.js';5import { operatorNotesOf, scriptedProvider, toolResultsOf, userTextsOf, type ScriptItem } from '../src/testkit/index.js';6import {7 cancelQuietly, createTestDb, getRun, makeWorkspace, newMission, pendingWait, sleep, startEngine, steps, waitFor, waitStatus,8 type MissionRef, type TestDb, type TestWorker,9} from './helpers.js';1011let db: TestDb;12let worker: TestWorker | null = null;13let mission: MissionRef | null = null;14beforeAll(async () => {15 db = await createTestDb();16});17afterEach(async () => {18 await worker?.stop();19 worker = null;20 if (mission) await cancelQuietly(db, mission.missionId);21 mission = null;22});23afterAll(async () => {24 await db.cleanup();25});26async function start(script: ScriptItem[], over: Record<string, unknown> = {}, deps: Parameters<typeof startEngine>[2] = {}) {27 const ws = makeWorkspace(db, { files: { 'a.txt': 'a\n' } });28 const provider = scriptedProvider(script);29 mission = await newMission(db, ws, over);30 worker = startEngine(db, provider, deps);31 return { provider, m: mission, ws };32}33const READ = { name: 'read_file', input: { path: 'a.txt', offset: null, limit: null } };34const LIST = (p: string) => ({ name: 'list_dir', input: { path: p, depth: null } });35const reads = (n: number): ScriptItem[] => Array.from({ length: n }, (_, i) => ({ text: `step ${i}`, calls: [LIST(i % 2 === 0 ? '.' : './')] }));36const PARTIAL = (summary: string) => ({ name: 'finish', input: { status: 'partial', summary, reason: 'budget exhausted', evidence: [], openQuestions: [] } });3738describe('budget governors', () => {39 it('turn limit: one final summarising turn, then budget_exhausted', async () => {40 const { provider, m } = await start([...reads(4), { text: 'Summary: I listed the directory four times; nothing is committed.', calls: [READ] }, ...reads(3)], { limits: { turns: 5 } });41 const run = await waitStatus(db, m.runId);42 expect(run.status).toBe('budget_exhausted');43 expect(provider.requests).toHaveLength(5); // never a 6th call44 expect(run.usage.turns).toBe(5);45 const notes = operatorNotesOf(provider.lastTranscript());46 expect(notes).toHaveLength(1);47 expect(notes[0]).toMatch(/turns: 4 of 5 used.*FINAL turn/s);48 expect(run.result?.summary).toMatch(/Budget exhausted — turns: 4 of 5 used/);49 expect(run.result?.summary).toContain('Summary: I listed the directory four times');50 // the tool call of the final turn was not executed51 expect((await steps(db, m.runId)).filter((s) => s.kind === 'tool_call')).toHaveLength(4);52 });5354 it('final turn: finish(partial) becomes budget_exhausted with the agent summary; finish(done) with passing criteria stays done', async () => {55 const a = await start([...reads(2), { calls: [PARTIAL('Half of the work is committed on the branch.')] }], { limits: { turns: 3 } });56 const ra = await waitStatus(db, a.m.runId);57 expect(ra.status).toBe('budget_exhausted');58 expect(ra.result?.summary).toMatch(/^Budget exhausted — turns: 2 of 3 used.*Half of the work is committed/s);59 await worker?.stop();60 const b = await start(61 [62 { calls: [{ name: 'update_criteria', input: { updates: [{ id: 'c1', passes: true, evidence: 'verified' }] } }] },63 { calls: [LIST('.')] },64 { calls: [{ name: 'finish', input: { status: 'done', summary: 'Finished in time.', reason: null, evidence: [], openQuestions: [] } }] },65 ],66 { limits: { turns: 3 } },67 );68 expect((await waitStatus(db, b.m.runId)).status).toBe('done');69 });7071 it('money: final turn at ≥ 90 % of the run budget', async () => {72 // each scripted turn costs USD 0.0031 → 93 % after three turns73 const provider = scriptedProvider([...reads(3), { text: 'Money summary.' }, ...reads(2)], { usage: { usd: 0.0031 } });74 const m = (mission = await newMission(db, makeWorkspace(db, { files: { 'a.txt': 'a\n' } }), { budgetUsd: 0.01 }));75 worker = startEngine(db, provider);76 expect(provider.requests.length).toBeLessThanOrEqual(4);77 const run = await waitStatus(db, m.runId);78 expect(run.status).toBe('budget_exhausted');79 expect(run.result?.summary).toMatch(/money: USD 0\.0093 of 0\.01 used \(93 %\)/);80 expect(run.result?.summary).toContain('Money summary.');81 expect(run.usage.turns).toBe(4);82 }, 60_000);8384 it('money: the kernel hard stop (no final turn possible) still ends budget_exhausted', async () => {85 const provider = scriptedProvider(reads(6), { usage: { usd: 0.006 } });86 const ws = makeWorkspace(db, { files: { 'a.txt': 'a\n' } });87 mission = await newMission(db, ws, { budgetUsd: 0.01 });88 worker = startEngine(db, provider);89 const run = await waitStatus(db, mission.runId);90 expect(run.status).toBe('budget_exhausted');91 expect(provider.requests).toHaveLength(2); // 0.006, 0.012 ≥ budget → the third call is refused by the kernel92 expect(run.result?.summary).toMatch(/Budget exhausted \(run_budget\)/);93 expect(run.result?.summary).toContain('step 1');94 });9596 it('wall clock: enforced by the engine, waits on the owner are not charged', async () => {97 const slow = (n: number): ScriptItem[] => Array.from({ length: n }, () => ({ calls: [LIST('.')], delayMs: 450 }));98 const { provider, m } = await start([...slow(2), { text: 'Out of time summary.' }, ...slow(3)], { limits: { wallSec: 1 } });99 const run = await waitStatus(db, m.runId);100 expect(run.status).toBe('budget_exhausted');101 expect(run.result?.summary).toMatch(/Budget exhausted — time: \d+ s of 1 s used/);102 expect(run.result?.summary).toContain('Out of time summary.');103 expect(provider.requests.length).toBeLessThanOrEqual(4);104 await worker?.stop();105 // a 1.5 s wait on the owner does not eat a 2 s wall budget106 const b = await start(107 [{ calls: [{ name: 'ask_owner', input: { question: 'Continue?', options: null } }] }, { calls: [LIST('.')] }, { calls: [PARTIAL('fine')] }],108 { limits: { wallSec: 2 } },109 );110 const w = await pendingWait(db, b.m.runId);111 await waitStatus(db, b.m.runId, ['waiting']);112 await sleep(2500);113 await resolveWait(w.id, { answer: 'yes' }, 'owner', db.k);114 const rb = await waitStatus(db, b.m.runId);115 expect(rb.status).toBe('partial');116 expect(operatorNotesOf(b.provider.lastTranscript())).toEqual([]);117 });118119 it('pure: budget view and final-turn rule', () => {120 const v = budgetView({ budget: { usd: 10, turns: 200, wallSec: 3600 }, left: { usd: 9, turns: 20, wallSec: 3000 }, usage: { usd: 1, inputTokens: 0, outputTokens: 0, cacheReadTokens: 0, cacheWriteTokens: 0, turns: 180, toolCalls: 0 }, waitedSec: 0 });121 expect(v).toMatchObject({ limit: 'turns', exhausted: false });122 expect(v.fraction).toBeCloseTo(0.9);123 expect(finalTurnDue(v, { turns: 20 }, 0.9)).toBe(true);124 const w = budgetView({ budget: { usd: 10, turns: 200, wallSec: 3600 }, left: { usd: 9, turns: 100, wallSec: -400 }, usage: { usd: 1, inputTokens: 0, outputTokens: 0, cacheReadTokens: 0, cacheWriteTokens: 0, turns: 100, toolCalls: 0 }, waitedSec: 3000 });125 expect(w.limit).toBe('turns'); // 4 000 s elapsed, 3 000 s of them waiting on the owner126 expect(finalTurnDue(w, { turns: 100 }, 0.9)).toBe(false);127 expect(finalTurnDue(w, { turns: 1 }, 0.9)).toBe(true);128 });129});130131describe('context guard', () => {132 it('warns at 70 % of the context window and ends partial at 90 %', async () => {133 const turn = (inputTokens: number): ScriptItem => ({ text: `ctx ${inputTokens}`, calls: [LIST(inputTokens % 200 === 100 ? '.' : './')], usage: { inputTokens, outputTokens: 10 } });134 const { provider, m } = await start([turn(300), turn(700), turn(800), turn(950), turn(990)], {}, { contextWindow: () => 1000 });135 const run = await waitStatus(db, m.runId);136 expect(run.status).toBe('partial');137 expect(provider.requests).toHaveLength(4);138 expect(run.result?.summary).toMatch(/fills about 96 % of the model's context window/);139 expect(run.result?.summary).toContain('ctx 950');140 const notes = operatorNotesOf(provider.lastTranscript());141 expect(notes).toHaveLength(1); // warned once, at 71 %142 expect(notes[0]).toMatch(/fills about 71 % of the model's context window/);143 });144});145146describe('stall detection', () => {147 it('same call + same result 3 times: an operator note, then a question to the owner', async () => {148 const same = (n: number): ScriptItem[] => Array.from({ length: n }, () => ({ text: 'Trying again.', calls: [READ] }));149 const { provider, m } = await start([...same(6), { calls: [PARTIAL('I stop.')] }]);150 const w = await pendingWait(db, m.runId);151 expect(w).toMatchObject({ kind: 'question', payload: { source: 'stall', options: ['Let it continue', 'Stop the mission'] } });152 expect(String(w.payload.detail)).toMatch(/same tool call returned the same result 3 times in a row/);153 expect(provider.requests).toHaveLength(6);154 expect(operatorNotesOf(provider.lastTranscript()).filter((n) => /You appear stuck/.test(n))).toHaveLength(1);155 await resolveWait(w.id, { answer: 'Read b.txt instead, it does not exist: just finish.' }, 'owner', db.k);156 expect((await waitStatus(db, m.runId)).status).toBe('partial');157 expect(userTextsOf(provider.lastTranscript()).at(-1)).toBe('Message from the owner:\nRead b.txt instead, it does not exist: just finish.');158 });159160 it('the owner can stop a stuck mission', async () => {161 const { m } = await start(Array.from({ length: 8 }, () => ({ calls: [READ] })));162 const w = await pendingWait(db, m.runId);163 await resolveWait(w.id, { answer: 'Stop the mission' }, 'owner', db.k);164 const run = await waitStatus(db, m.runId);165 expect(run.status).toBe('partial');166 expect(run.result?.summary).toMatch(/Stopped by the owner after the agent appeared stuck/);167 });168169 it('N turns without a write, commit or criterion change', async () => {170 const { provider, m } = await start(171 [...reads(3), { calls: [{ name: 'write_file', input: { path: 'n.txt', content: 'n' } }] }, ...reads(2), { calls: [PARTIAL('ok')] }],172 {},173 { limits: { stallTurns: 3 } },174 );175 expect((await waitStatus(db, m.runId)).status).toBe('partial');176 const notes = operatorNotesOf(provider.lastTranscript()).filter((n) => /You appear stuck/.test(n));177 expect(notes).toHaveLength(1); // after turns 1–3; the write at turn 4 reset the counter178 expect(notes[0]).toMatch(/3 turns have passed without a successful write/);179 expect(provider.requests[3]?.transcript.filter((e) => e.role === 'operator')).toHaveLength(1);180 });181182 it('pure: detector', () => {183 const d = new StallDetector(3, 100);184 d.recordCall('t', { a: 1, b: 2 }, true, 'x');185 d.recordCall('t', { b: 2, a: 1 }, true, 'x'); // canonical input: same call186 expect(d.check()).toBe('none');187 d.recordCall('t', { a: 1, b: 2 }, true, 'y'); // different result: the streak restarts188 d.recordCall('t', { a: 1, b: 2 }, true, 'y');189 expect(d.check()).toBe('none');190 d.recordCall('t', { a: 1, b: 2 }, true, 'y');191 expect(d.check()).toBe('note');192 for (let i = 0; i < 3; i++) d.recordCall('t', {}, false, 'e');193 expect(d.check()).toBe('ask');194 for (let i = 0; i < 3; i++) d.recordCall('t', {}, false, 'e');195 expect(d.check()).toBe('note');196 });197});198199describe('cancel and steering', () => {200 it('cancel while a model call is in flight ends the run cancelled', async () => {201 const { m } = await start([{ calls: [LIST('.')] }, { text: 'thinking for a long time', delayMs: 20_000 }]);202 await waitFor(async () => (await steps(db, m.runId)).filter((s) => s.kind === 'model_call').length === 2, 20_000, 'second model call');203 await requestCancel(m.missionId, db.k);204 const run = await waitStatus(db, m.runId, undefined, 10_000);205 expect(run.status).toBe('cancelled');206 expect((await db.k.pool.query('SELECT status FROM missions WHERE id = $1', [m.missionId])).rows[0]).toEqual({ status: 'cancelled' });207 });208209 it('cancel while waiting for an approval', async () => {210 const { m } = await start([{ calls: [{ name: 'ask_owner', input: { question: 'Go?', options: null } }] }]);211 await pendingWait(db, m.runId);212 await waitStatus(db, m.runId, ['waiting']);213 await requestCancel(m.missionId, db.k);214 expect((await getRun(db, m.runId)).status).toBe('cancelled');215 });216217 it('a steering message sent mid-run is read right after the current tool batch', async () => {218 let posted = false;219 const { provider, m } = await start([220 { calls: [LIST('.')], delayMs: 300 },221 { calls: [LIST('.')] },222 { calls: [PARTIAL('ok')] },223 ]);224 await waitFor(async () => (await steps(db, m.runId)).some((s) => s.kind === 'model_call' && s.status === 'started'), 10_000, 'first model call');225 await postOwnerMessage(m.runId, 'Also check the README.', db.k);226 posted = true;227 expect((await waitStatus(db, m.runId)).status).toBe('partial');228 expect(posted).toBe(true);229 // turn 2 already sees it, after the tool results of turn 1230 const t2 = provider.requests[1]?.transcript ?? [];231 expect(t2.map((e) => e.role)).toEqual(['user', 'assistant', 'tool', 'user']);232 expect(userTextsOf(t2)[1]).toBe('Message from the owner:\nAlso check the README.');233 expect(toolResultsOf(t2)).toHaveLength(1);234 });235});236