File size: 5,458 Bytes
4e23b01 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import type { Event } from '../src/transport/ws/v1/events';
import { describe, expect, it } from 'vitest';
import { InFlightTurnTracker } from '../src/transport/ws/v1/inFlightTurnTracker';
const SID = 'sess_1';
function ev(partial: Record<string, unknown>): Event {
return { agentId: 'main', sessionId: SID, ...partial } as unknown as Event;
}
describe('InFlightTurnTracker', () => {
it('accumulates assistant text and reports pre-append offsets', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
expect(t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'Hello' }))).toEqual({
offset: 0,
});
expect(t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: ' world' }))).toEqual({
offset: 5,
});
expect(t.get(SID)).toMatchObject({ turn_id: 1, assistant_text: 'Hello world' });
});
it('tracks thinking offsets independently', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
expect(t.apply(SID, ev({ type: 'thinking.delta', turnId: 1, delta: 'abc' }))).toEqual({
offset: 0,
});
expect(t.apply(SID, ev({ type: 'thinking.delta', turnId: 1, delta: 'de' }))).toEqual({
offset: 3,
});
expect(t.get(SID)).toMatchObject({ assistant_text: '', thinking_text: 'abcde' });
});
it('clears on turn.ended', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'x' }));
t.apply(SID, ev({ type: 'turn.ended', turnId: 1 }));
expect(t.get(SID)).toBeNull();
});
it('ignores non-main agents', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
const sub = { agentId: 'agent-sub', sessionId: SID, type: 'assistant.delta', turnId: 1, delta: 'nope' } as unknown as Event;
expect(t.apply(SID, sub)).toEqual({});
expect(t.get(SID)?.assistant_text).toBe('');
});
it('ignores deltas for a mismatched turn', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
expect(t.apply(SID, ev({ type: 'assistant.delta', turnId: 99, delta: 'stale' }))).toEqual({});
expect(t.get(SID)?.assistant_text).toBe('');
});
it('tracks running tools and their last progress', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
t.apply(SID, ev({ type: 'tool.call.started', turnId: 1, toolCallId: 'tc1', name: 'bash' }));
t.apply(
SID,
ev({ type: 'tool.progress', turnId: 1, toolCallId: 'tc1', update: { kind: 'stdout', text: 'hi' } }),
);
expect(t.get(SID)?.running_tools).toEqual([
{ tool_call_id: 'tc1', name: 'bash', last_progress: { kind: 'stdout', text: 'hi' } },
]);
t.apply(SID, ev({ type: 'tool.result', turnId: 1, toolCallId: 'tc1' }));
expect(t.get(SID)?.running_tools).toEqual([]);
});
it('resets text accumulation at step boundaries (step-relative in-flight text)', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
t.apply(SID, ev({ type: 'turn.step.started', turnId: 1, step: 1 }));
t.apply(SID, ev({ type: 'thinking.delta', turnId: 1, delta: 'think-1' }));
t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'text-1' }));
t.apply(SID, ev({ type: 'turn.step.completed', turnId: 1, step: 1 }));
t.apply(SID, ev({ type: 'turn.step.started', turnId: 1, step: 2 }));
t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'text-2' }));
expect(t.get(SID)).toMatchObject({ assistant_text: 'text-2', thinking_text: '' });
});
it('reports step-relative offsets that restart at 0 each step', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
t.apply(SID, ev({ type: 'turn.step.started', turnId: 1, step: 1 }));
expect(t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'ab' }))).toEqual({ offset: 0 });
expect(t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'cd' }))).toEqual({ offset: 2 });
t.apply(SID, ev({ type: 'turn.step.started', turnId: 1, step: 2 }));
expect(t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'x' }))).toEqual({ offset: 0 });
});
it('keeps running tools across step boundaries while resetting text', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
t.apply(SID, ev({ type: 'turn.step.started', turnId: 1, step: 1 }));
t.apply(SID, ev({ type: 'tool.call.started', turnId: 1, toolCallId: 'tc1', name: 'bash' }));
t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'text-1' }));
t.apply(SID, ev({ type: 'turn.step.started', turnId: 1, step: 2 }));
expect(t.get(SID)?.assistant_text).toBe('');
expect(t.get(SID)?.running_tools).toEqual([{ tool_call_id: 'tc1', name: 'bash' }]);
});
it('ignores step boundaries for a mismatched turn', () => {
const t = new InFlightTurnTracker();
t.apply(SID, ev({ type: 'turn.started', turnId: 1 }));
t.apply(SID, ev({ type: 'assistant.delta', turnId: 1, delta: 'keep' }));
t.apply(SID, ev({ type: 'turn.step.started', turnId: 99, step: 2 }));
expect(t.get(SID)?.assistant_text).toBe('keep');
});
});
|