Skip to content

Commit 2e11a74

Browse files
bloveclaude
blove
and
claude
committed
feat(langgraph): make MockAgentTransport emit/emitError/close awaitable, add flush()
`stream()` is an async generator, so `emit()` only woke the suspended loop and nothing had reached the signals when it returned. Every spec paid for that with a hand-rolled `await new Promise(r => setTimeout(r, 0))`. `emit()`, `emitError()` and `close()` now return a promise that settles once the generator has drained everything queued at the time of the call (or the run has ended), plus one macrotask so signal writes have landed. `flush()` waits the same way without emitting. An emit after the run finished resolves instead of hanging. The langgraph specs that hand-rolled the macrotask flush after an emit are converted to `await transport.emit(...)`; removing the await makes them fail, so the await is load-bearing. Throttle waits (16 ms and up) are left alone. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 parent b0b1e0b commit 2e11a74

4 files changed

Lines changed: 177 additions & 44 deletions

File tree

‎libs/langgraph/src/lib/agent.fn.spec.ts‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -133,12 +133,11 @@ describe('agent', () => {
133133
);
134134

135135
const submitted = ref.submit({ message: 'hello' });
136-
transport.emit([{
136+
await transport.emit([{
137137
type: 'messages',
138138
messages: [{ id: 'ai-live', type: 'ai', content: 'answer' }],
139139
messageMetadata: { langgraph_node: 'model' },
140140
}]);
141-
await new Promise(resolve => setTimeout(resolve, 0));
142141

143142
const streaming = ref.messages().find(message => message.id === 'ai-live')?.delivery;
144143
expect(streaming).toEqual({ generation: expect.any(String), phase: 'streaming' });
@@ -206,7 +205,7 @@ describe('agent', () => {
206205
],
207206
}],
208207
}]);
209-
transport.emit([{
208+
await transport.emit([{
210209
type: 'messages|tools:call-success', namespace: ['tools:call-success'],
211210
messages: [{ id: 'sub-success', type: 'ai', content: 'result' }],
212211
messageMetadata: { checkpoint_ns: 'tools:call-success|model' },
@@ -215,22 +214,20 @@ describe('agent', () => {
215214
messages: [{ id: 'sub-error', type: 'ai', content: 'partial' }],
216215
messageMetadata: { checkpoint_ns: 'tools:call-error|model' },
217216
}]);
218-
await new Promise(resolve => setTimeout(resolve, 0));
219217

220218
const successStreaming = ref.subagents().get('call-success')?.messages()[0].delivery;
221219
const errorStreaming = ref.subagents().get('call-error')?.messages()[0].delivery;
222220
expect(successStreaming).toMatchObject({ phase: 'streaming' });
223221
expect(errorStreaming).toMatchObject({ phase: 'streaming' });
224222
expect(successStreaming?.generation).not.toBe(errorStreaming?.generation);
225223

226-
transport.emit([{
224+
await transport.emit([{
227225
type: 'messages',
228226
messages: [
229227
{ id: 'tool-success', type: 'tool', tool_call_id: 'call-success', content: 'done', status: 'success' },
230228
{ id: 'tool-error', type: 'tool', tool_call_id: 'call-error', content: 'failed', status: 'error' },
231229
],
232230
}]);
233-
await new Promise(resolve => setTimeout(resolve, 0));
234231

235232
expect(ref.subagents().get('call-success')?.messages()[0].delivery).toEqual({
236233
generation: successStreaming?.generation,

‎libs/langgraph/src/lib/internals/stream-manager.bridge.spec.ts‎

Lines changed: 13 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -374,12 +374,11 @@ describe('createStreamManagerBridge', () => {
374374
});
375375

376376
const submitted = bridge.submit({});
377-
transport.emit([{
377+
await transport.emit([{
378378
type: 'messages',
379379
messages: [{ id: 'ai-1', type: 'ai', content: 'hel' }],
380380
messageMetadata: { langgraph_node: 'model' },
381381
}]);
382-
await new Promise(resolve => setTimeout(resolve, 0));
383382

384383
const streaming = bridge.getMessageDelivery('ai-1');
385384
expect(streaming).toEqual({
@@ -411,20 +410,18 @@ describe('createStreamManagerBridge', () => {
411410
});
412411

413412
void bridge.submit({});
414-
transport.emit([{
413+
await transport.emit([{
415414
type: 'messages',
416415
messages: [{ id: 'revision-ai', type: 'ai', content: 'a' }],
417416
messageMetadata: { langgraph_node: 'model' },
418417
}]);
419-
await new Promise(resolve => setTimeout(resolve, 0));
420418
const afterFirstChunk = bridge.deliveryRevision();
421419

422-
transport.emit([{
420+
await transport.emit([{
423421
type: 'messages',
424422
messages: [{ id: 'revision-ai', type: 'ai', content: 'b' }],
425423
messageMetadata: { langgraph_node: 'model' },
426424
}]);
427-
await new Promise(resolve => setTimeout(resolve, 0));
428425

429426
expect(bridge.deliveryRevision()).toBe(afterFirstChunk);
430427
await bridge.stop();
@@ -444,12 +441,11 @@ describe('createStreamManagerBridge', () => {
444441
await new Promise(resolve => setTimeout(resolve, 0));
445442

446443
const submitted = bridge.submit({});
447-
transport.emit([{
444+
await transport.emit([{
448445
type: 'messages',
449446
messages: [{ id: 'streamed-id', type: 'ai', content: 'final answer' }],
450447
messageMetadata: { langgraph_node: 'model' },
451448
}]);
452-
await new Promise(resolve => setTimeout(resolve, 0));
453449
const streaming = bridge.getMessageDelivery('streamed-id');
454450

455451
transport.history = [{
@@ -998,12 +994,11 @@ describe('createStreamManagerBridge', () => {
998994
});
999995

1000996
const submitted = bridge.submit({});
1001-
transport.emit([{
997+
await transport.emit([{
1002998
type: 'messages',
1003999
messages: [{ id: 'ai-aborted', type: 'ai', content: 'partial' }],
10041000
messageMetadata: { langgraph_node: 'model' },
10051001
}]);
1006-
await new Promise(resolve => setTimeout(resolve, 0));
10071002
await bridge.stop();
10081003
transport.close();
10091004

@@ -1397,21 +1392,19 @@ describe('createStreamManagerBridge', () => {
13971392
});
13981393

13991394
void bridge.submit({});
1400-
transport.emit([{
1395+
await transport.emit([{
14011396
type: 'messages',
14021397
messages: [{ id: 'ai-tool-call', type: 'ai', content: 'search', tool_calls: [{ id: 'call-1', name: 'search', args: {} }] }],
14031398
messageMetadata: { langgraph_node: 'model' },
14041399
}]);
1405-
await new Promise(resolve => setTimeout(resolve, 0));
14061400
expect(bridge.getMessageDelivery('ai-tool-call').phase).toBe('streaming');
14071401

14081402
transport.emit([{ type: 'values', values: { toolStepComplete: true } }]);
1409-
transport.emit([{
1403+
await transport.emit([{
14101404
type: 'messages',
14111405
messages: [{ id: 'ai-final', type: 'ai', content: 'search complete' }],
14121406
messageMetadata: { langgraph_node: 'model' },
14131407
}]);
1414-
await new Promise(resolve => setTimeout(resolve, 0));
14151408

14161409
expect(bridge.getMessageDelivery('ai-tool-call')).toMatchObject({
14171410
phase: 'complete',
@@ -1434,7 +1427,7 @@ describe('createStreamManagerBridge', () => {
14341427
});
14351428

14361429
void bridge.submit({});
1437-
transport.emit([{
1430+
await transport.emit([{
14381431
type: 'messages',
14391432
messages: [{
14401433
id: 'tool-chunk-a', type: 'ai', content: 'hel',
@@ -1449,7 +1442,6 @@ describe('createStreamManagerBridge', () => {
14491442
}],
14501443
messageMetadata: { langgraph_node: 'model' },
14511444
}]);
1452-
await new Promise(resolve => setTimeout(resolve, 0));
14531445

14541446
expect(subjects.messages$.value).toEqual([
14551447
expect.objectContaining({
@@ -1484,7 +1476,7 @@ describe('createStreamManagerBridge', () => {
14841476
});
14851477

14861478
const submitted = bridge.submit({});
1487-
transport.emit([{
1479+
await transport.emit([{
14881480
type: 'messages',
14891481
messages: [{ id: 'chunk-event-1', type: 'ai', content: 'hel' }],
14901482
messageMetadata: { langgraph_node: 'model' },
@@ -1493,7 +1485,6 @@ describe('createStreamManagerBridge', () => {
14931485
messages: [{ id: 'chunk-event-2', type: 'ai', content: 'lo' }],
14941486
messageMetadata: { langgraph_node: 'model' },
14951487
}]);
1496-
await new Promise(resolve => setTimeout(resolve, 0));
14971488

14981489
expect(subjects.messages$.value).toEqual([
14991490
expect.objectContaining({ id: 'chunk-event-1', content: 'hello' }),
@@ -1535,7 +1526,7 @@ describe('createStreamManagerBridge', () => {
15351526
});
15361527

15371528
const submitted = bridge.submit({});
1538-
transport.emit([{
1529+
await transport.emit([{
15391530
type: 'messages',
15401531
messages: [{
15411532
id: 'ai-earlier', type: 'ai', content: '',
@@ -1553,7 +1544,6 @@ describe('createStreamManagerBridge', () => {
15531544
messages: [{ id: 'ai-active', type: 'ai', content: 'partial' }],
15541545
messageMetadata: { langgraph_node: 'model' },
15551546
}]);
1556-
await new Promise(resolve => setTimeout(resolve, 0));
15571547

15581548
if (outcome === 'error') {
15591549
transport.emit([{ type: 'error', error: new Error('failed') }]);
@@ -2508,15 +2498,14 @@ describe('createStreamManagerBridge', () => {
25082498
});
25092499

25102500
void bridge.submit({});
2511-
transport.emit([{
2501+
await transport.emit([{
25122502
type,
25132503
messages: [
25142504
{ id: 'historical-ai', type: 'ai', content: 'old answer' },
25152505
{ id: 'historical-user', type: 'human', content: 'new question' },
25162506
{ id: 'active-ai', type: 'ai', content: 'new answer' },
25172507
],
25182508
}]);
2519-
await new Promise(resolve => setTimeout(resolve, 0));
25202509

25212510
expect(bridge.getMessageDelivery('historical-ai')).toEqual({
25222511
generation: 'historical-ai',
@@ -2551,7 +2540,7 @@ describe('createStreamManagerBridge', () => {
25512540
const historicalDelivery = bridge.getMessageDelivery('historical-enriched-ai');
25522541

25532542
void bridge.submit({});
2554-
transport.emit([{
2543+
await transport.emit([{
25552544
type,
25562545
messages: [
25572546
{
@@ -2565,7 +2554,6 @@ describe('createStreamManagerBridge', () => {
25652554
{ id: 'active-enriched-ai', type: 'ai', content: 'new answer' },
25662555
],
25672556
}]);
2568-
await new Promise(resolve => setTimeout(resolve, 0));
25692557

25702558
expect(subjects.messages$.value.find(message =>
25712559
(message as unknown as { id?: string }).id === 'historical-enriched-ai'
@@ -3525,13 +3513,12 @@ describe('createStreamManagerBridge', () => {
35253513
});
35263514

35273515
const done = bridge.submit({});
3528-
transport.emit([{
3516+
await transport.emit([{
35293517
type: 'messages|research:abc123' as StreamEvent['type'],
35303518
namespace: ['research:abc123'],
35313519
messages: [{ id: 'child-ai', type: 'ai', content: 'brief' }],
35323520
messageMetadata: { checkpoint_ns: 'research:abc123' },
35333521
} satisfies StreamEvent]);
3534-
await new Promise(r => setTimeout(r, 0));
35353522
expect(subjects.subagents$.value.get('research:abc123')?.status()).toBe('running');
35363523

35373524
transport.emit([{ type: 'values', data: { done: true } } as StreamEvent]);

‎libs/langgraph/src/lib/transport/mock-stream.transport.spec.ts‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { describe, it, expect, beforeEach } from 'vitest';
1+
import { describe, it, expect } from 'vitest';
22
import { MockAgentTransport } from './mock-stream.transport';
33

44
describe('MockAgentTransport', () => {
@@ -39,6 +39,91 @@ describe('MockAgentTransport', () => {
3939
expect(events).toHaveLength(1);
4040
});
4141

42+
describe('awaitable emit()', () => {
43+
it('emit() resolves only after the consumer has pulled the batch', async () => {
44+
const t = new MockAgentTransport();
45+
const events: unknown[] = [];
46+
const ac = new AbortController();
47+
const collecting = (async () => {
48+
for await (const e of t.stream('agent', null, {}, ac.signal)) { events.push(e); }
49+
})();
50+
51+
await t.emit([
52+
{ type: 'values', values: { foo: 1 } },
53+
{ type: 'values', values: { foo: 2 } },
54+
]);
55+
// No bare setTimeout: awaiting emit() is enough for the batch to land.
56+
expect(events).toHaveLength(2);
57+
58+
await t.close();
59+
await collecting;
60+
});
61+
62+
it('emit() before the stream starts resolves once the stream drains it', async () => {
63+
const t = new MockAgentTransport();
64+
const events: unknown[] = [];
65+
const ac = new AbortController();
66+
const emitted = t.emit([{ type: 'values', values: { foo: 1 } }]);
67+
const collecting = (async () => {
68+
for await (const e of t.stream('agent', null, {}, ac.signal)) { events.push(e); }
69+
})();
70+
await emitted;
71+
expect(events).toHaveLength(1);
72+
await t.close();
73+
await collecting;
74+
});
75+
76+
it('flush() resolves without emitting anything', async () => {
77+
const t = new MockAgentTransport();
78+
const ac = new AbortController();
79+
const collecting = (async () => {
80+
for await (const _ of t.stream('agent', null, {}, ac.signal)) { /* noop */ }
81+
})();
82+
await t.emit([{ type: 'values', values: { foo: 1 } }]);
83+
await t.flush();
84+
expect(t.isStreaming()).toBe(true);
85+
await t.close();
86+
await collecting;
87+
});
88+
89+
it('close() resolves once the run has finished', async () => {
90+
const t = new MockAgentTransport();
91+
const ac = new AbortController();
92+
const collecting = (async () => {
93+
for await (const _ of t.stream('agent', null, {}, ac.signal)) { /* noop */ }
94+
})();
95+
await t.close();
96+
expect(t.isStreaming()).toBe(false);
97+
await collecting;
98+
});
99+
100+
it('emitError() resolves once the stream has thrown', async () => {
101+
const t = new MockAgentTransport();
102+
const ac = new AbortController();
103+
let thrown: unknown;
104+
const collecting = (async () => {
105+
try {
106+
for await (const _ of t.stream('agent', null, {}, ac.signal)) { /* noop */ }
107+
} catch (e) { thrown = e; }
108+
})();
109+
await t.emitError(new Error('transport error'));
110+
expect(thrown).toBeInstanceOf(Error);
111+
await collecting;
112+
});
113+
114+
it('emit() after the run has finished resolves instead of hanging', async () => {
115+
const t = new MockAgentTransport();
116+
const ac = new AbortController();
117+
const collecting = (async () => {
118+
for await (const _ of t.stream('agent', null, {}, ac.signal)) { /* noop */ }
119+
})();
120+
await t.close();
121+
await collecting;
122+
await expect(t.emit([{ type: 'values', values: { foo: 1 } }])).resolves.toBeUndefined();
123+
await expect(t.flush()).resolves.toBeUndefined();
124+
});
125+
});
126+
42127
it('emitError() causes stream to throw', async () => {
43128
const t = new MockAgentTransport();
44129
const ac = new AbortController();

0 commit comments

Comments
 (0)