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');
  });
});