117 lines
4.4 KiB
TypeScript
117 lines
4.4 KiB
TypeScript
import { usageOf } from './helpers';
|
|
import { expect, test } from 'bun:test';
|
|
import { MockLanguageModelV4, simulateReadableStream } from 'ai/test';
|
|
import type { LanguageModelV4StreamPart } from '@ai-sdk/provider';
|
|
import { Session, type AgentEvent } from '../src/session';
|
|
|
|
function stream(parts: LanguageModelV4StreamPart[]) {
|
|
return { stream: simulateReadableStream({ chunks: parts, chunkDelayInMs: null, initialDelayInMs: null }) };
|
|
}
|
|
|
|
/** A priced model whose single run bills `usage`, so the ceiling arithmetic is exact. */
|
|
function modelWith(usage: ReturnType<typeof usageOf>) {
|
|
return new MockLanguageModelV4({
|
|
doStream: async () =>
|
|
stream([
|
|
{ type: 'text-start', id: '0' },
|
|
{ type: 'text-delta', id: '0', delta: 'ok' },
|
|
{ type: 'text-end', id: '0' },
|
|
{ type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage },
|
|
]),
|
|
});
|
|
}
|
|
|
|
async function drain(session: Session): Promise<AgentEvent[]> {
|
|
const events: AgentEvent[] = [];
|
|
for await (const ev of session.send('go')) events.push(ev);
|
|
return events;
|
|
}
|
|
|
|
const noop = async () => 'once' as const;
|
|
|
|
test('a session past its ceiling refuses the turn and names the ceiling', async () => {
|
|
// gpt-5 is $1.25/M in, $10/M out. 800k in / 200k out = $1.00 + $2.00 = $3.00.
|
|
const session = new Session({
|
|
model: modelWith(usageOf(800_000, 200_000)),
|
|
modelId: 'gpt-5',
|
|
askApproval: noop,
|
|
maxSpendUsd: 2,
|
|
});
|
|
|
|
const first = await drain(session);
|
|
expect(first.some((e) => e.type === 'done')).toBe(true);
|
|
|
|
// The first turn put spend over the $2 ceiling; the next is refused before the model runs.
|
|
const second = await drain(session);
|
|
const refusal = second.find((e) => e.type === 'error');
|
|
expect(refusal).toBeDefined();
|
|
expect(refusal!.type === 'error' && String(refusal!.error)).toContain('spend ceiling reached');
|
|
expect(String(refusal!.type === 'error' && refusal!.error)).toContain('$2.00');
|
|
});
|
|
|
|
test('crossing 80% warns once, then stays quiet', async () => {
|
|
// 200k in / 100k out = $0.25 + $1.00 = $1.25, which is 83% of a $1.50 ceiling.
|
|
const session = new Session({
|
|
model: modelWith(usageOf(200_000, 100_000)),
|
|
modelId: 'gpt-5',
|
|
askApproval: noop,
|
|
maxSpendUsd: 1.5,
|
|
});
|
|
|
|
const first = await drain(session);
|
|
const warnings = first.filter((e) => e.type === 'notice' && e.text.includes('approaching spend ceiling'));
|
|
expect(warnings).toHaveLength(1);
|
|
expect(warnings[0]!.type === 'notice' && warnings[0]!.text).toContain('$1.50');
|
|
|
|
// A second turn at the same level must not repeat the warning.
|
|
const second = await drain(session);
|
|
expect(second.filter((e) => e.type === 'notice' && e.text.includes('approaching spend ceiling'))).toHaveLength(0);
|
|
});
|
|
|
|
test('an unpriced model is never refused, because the ceiling cannot see it', async () => {
|
|
const session = new Session({
|
|
model: modelWith(usageOf(9_000_000, 9_000_000)),
|
|
modelId: 'some-local-model',
|
|
askApproval: noop,
|
|
maxSpendUsd: 0.01,
|
|
});
|
|
const events = await drain(session);
|
|
expect(events.some((e) => e.type === 'done')).toBe(true);
|
|
expect(events.some((e) => e.type === 'error')).toBe(false);
|
|
});
|
|
|
|
test('with no ceiling configured a session spends freely', async () => {
|
|
const session = new Session({ model: modelWith(usageOf(9_000_000, 9_000_000)), modelId: 'gpt-5', askApproval: noop });
|
|
const events = await drain(session);
|
|
expect(events.some((e) => e.type === 'done')).toBe(true);
|
|
});
|
|
|
|
test('spend() reports the ceiling state for the status surface', () => {
|
|
const session = new Session({ model: modelWith(usageOf(1)), modelId: 'gpt-5', askApproval: noop, maxSpendUsd: 4 });
|
|
session.inputTokens = 800_000;
|
|
session.outputTokens = 200_000; // $3.00 of $4.00
|
|
const s = session.spend();
|
|
expect(s.usd).toBeCloseTo(3, 5);
|
|
expect(s.ceiling).toBe(4);
|
|
expect(s.overWarn).toBe(false);
|
|
expect(s.overLimit).toBe(false);
|
|
});
|
|
|
|
test('subagent spend folds into the ceiling and resets with the session', async () => {
|
|
const session = new Session({
|
|
model: modelWith(usageOf(1)),
|
|
modelId: 'gpt-5',
|
|
subagentModelId: 'gpt-5-nano',
|
|
askApproval: noop,
|
|
maxSpendUsd: 4,
|
|
});
|
|
session.recordSubagentUsage({ inputTokens: 1_000_000, outputTokens: 100_000 });
|
|
expect(session.subagentInputTokens).toBe(1_000_000);
|
|
// gpt-5-nano: $0.05/M in, $0.40/M out -> $0.05 + $0.04 = $0.09
|
|
expect(session.spend().usd).toBeCloseTo(0.09, 5);
|
|
|
|
session.reset();
|
|
expect(session.subagentInputTokens).toBe(0);
|
|
expect(session.spend().usd).toBeCloseTo(0, 5);
|
|
});
|