Skip to content

Commit 1bf2aa5

Browse files
yyyyaaayyyyaaa
committed
feat(react): expose cumulative usage in useChat with change-detection guard
Surface the cumulative Usage from agent_end/turn_end as a top-level useChat field, null before the first event. Guard the setter on totalTokens so re-render does not fire when nothing changed. Reset to null on every new prompt(), mirroring the agent-side reset.
1 parent 0b40656 commit 1bf2aa5

2 files changed

Lines changed: 206 additions & 83 deletions

File tree

‎packages/react/__tests__/use-chat.test.ts‎

Lines changed: 122 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { AgentEvent } from '@agentic-kit/agent';
2-
import { createScriptedSSEResponse, makeFakeAssistantMessage } from '@test/index';
2+
import { createScriptedSSEResponse, makeFakeAssistantMessage, ZERO_USAGE } from '@test/index';
33
import { act, renderHook, waitFor } from '@testing-library/react';
44
import type { AssistantMessage, Message, UserMessage } from 'agentic-kit';
55

@@ -77,7 +77,7 @@ describe('useChat', () => {
7777
async (_url: RequestInfo | URL, _init?: RequestInit): Promise<Response> =>
7878
streamFromEvents([
7979
{ type: 'agent_start' },
80-
{ type: 'agent_end', messages: [makeUser('first'), makeUser('second'), final] },
80+
{ type: 'agent_end', messages: [makeUser('first'), makeUser('second'), final], totalUsage: ZERO_USAGE },
8181
])
8282
);
8383

@@ -118,7 +118,7 @@ describe('useChat', () => {
118118
},
119119
},
120120
{ type: 'message_end', message: final },
121-
{ type: 'agent_end', messages: [userEcho, final] },
121+
{ type: 'agent_end', messages: [userEcho, final], totalUsage: ZERO_USAGE },
122122
])
123123
);
124124
const onMessage = jest.fn();
@@ -196,6 +196,7 @@ describe('useChat', () => {
196196
pushFn({
197197
type: 'agent_end',
198198
messages: [makeUser('hi'), final],
199+
totalUsage: ZERO_USAGE,
199200
});
200201
closeFn();
201202
await act(async () => {
@@ -215,7 +216,7 @@ describe('useChat', () => {
215216
async (_url: RequestInfo | URL, _init?: RequestInit): Promise<Response> =>
216217
streamFromEvents([
217218
{ type: 'agent_start' },
218-
{ type: 'agent_end', messages: [makeUser('hi'), final] },
219+
{ type: 'agent_end', messages: [makeUser('hi'), final], totalUsage: ZERO_USAGE },
219220
])
220221
);
221222
const body = jest.fn(() => ({ model: 'demo', sessionId: 'abc' }));
@@ -248,7 +249,7 @@ describe('useChat', () => {
248249
controller.enqueue(encoder.encode('data: {garbage not json\n\n'));
249250
controller.enqueue(
250251
encoder.encode(
251-
`data: ${JSON.stringify({ type: 'agent_end', messages: [makeUser('hi'), final] })}\n\n`
252+
`data: ${JSON.stringify({ type: 'agent_end', messages: [makeUser('hi'), final], totalUsage: ZERO_USAGE })}\n\n`
252253
)
253254
);
254255
controller.close();
@@ -280,7 +281,7 @@ describe('useChat', () => {
280281
async (_url: RequestInfo | URL, _init?: RequestInit): Promise<Response> =>
281282
streamFromEvents([
282283
{ type: 'agent_start' },
283-
{ type: 'agent_end', messages: [makeUser('a'), makeUser('b'), final] },
284+
{ type: 'agent_end', messages: [makeUser('a'), makeUser('b'), final], totalUsage: ZERO_USAGE },
284285
])
285286
);
286287
const { result } = renderHook(() => useChat({ api: '/chat', fetch: fetchFn }));
@@ -304,7 +305,7 @@ describe('useChat', () => {
304305
});
305306
return streamFromEvents([
306307
{ type: 'agent_start' },
307-
{ type: 'agent_end', messages: [makeUser('hi'), makeFinalAssistant('hello')] },
308+
{ type: 'agent_end', messages: [makeUser('hi'), makeFinalAssistant('hello')], totalUsage: ZERO_USAGE },
308309
]);
309310
}
310311
);
@@ -430,7 +431,7 @@ describe('useChat', () => {
430431
isError: false,
431432
},
432433
{ type: 'message_end', message: final },
433-
{ type: 'agent_end', messages: [userEcho, final] },
434+
{ type: 'agent_end', messages: [userEcho, final], totalUsage: ZERO_USAGE },
434435
])
435436
);
436437

@@ -762,7 +763,7 @@ describe('useChat', () => {
762763
expect(result.current.isStreaming).toBe(false);
763764

764765
const lateAssistant = makeFinalAssistant('late');
765-
pushFn({ type: 'agent_end', messages: [makeUser('hi'), lateAssistant] });
766+
pushFn({ type: 'agent_end', messages: [makeUser('hi'), lateAssistant], totalUsage: ZERO_USAGE });
766767
closeFn();
767768
await act(async () => {
768769
await sendPromise;
@@ -843,6 +844,7 @@ describe('useChat', () => {
843844
{
844845
type: 'agent_end',
845846
messages: [userEcho, resumedAssistant, toolResult, final],
847+
totalUsage: ZERO_USAGE,
846848
},
847849
])
848850
);
@@ -875,7 +877,7 @@ describe('useChat', () => {
875877
async (): Promise<Response> =>
876878
streamFromEvents([
877879
{ type: 'agent_start' },
878-
{ type: 'agent_end', messages: initial },
880+
{ type: 'agent_end', messages: initial, totalUsage: ZERO_USAGE },
879881
])
880882
);
881883

@@ -907,7 +909,7 @@ describe('useChat', () => {
907909
async (_url: RequestInfo | URL, _init?: RequestInit): Promise<Response> =>
908910
streamFromEvents([
909911
{ type: 'agent_start' },
910-
{ type: 'agent_end', messages: initial },
912+
{ type: 'agent_end', messages: initial, totalUsage: ZERO_USAGE },
911913
])
912914
);
913915

@@ -964,7 +966,7 @@ describe('useChat', () => {
964966
async (_url: RequestInfo | URL, _init?: RequestInit): Promise<Response> =>
965967
streamFromEvents([
966968
{ type: 'agent_start' },
967-
{ type: 'agent_end', messages: initial },
969+
{ type: 'agent_end', messages: initial, totalUsage: ZERO_USAGE },
968970
])
969971
);
970972

@@ -1049,6 +1051,113 @@ describe('useChat', () => {
10491051
});
10501052
});
10511053

1054+
describe('usage', () => {
1055+
const sampleUsage = {
1056+
input: 10,
1057+
output: 20,
1058+
reasoning: 0,
1059+
cacheRead: 0,
1060+
cacheWrite: 0,
1061+
totalTokens: 30,
1062+
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
1063+
};
1064+
1065+
it('starts as null', () => {
1066+
const { result } = renderHook(() => useChat({ api: '/chat' }));
1067+
expect(result.current.usage).toBeNull();
1068+
});
1069+
1070+
it('is populated from agent_end totalUsage', async () => {
1071+
const final = makeFinalAssistant('ok');
1072+
const fetchFn = jest.fn(
1073+
async (): Promise<Response> =>
1074+
streamFromEvents([
1075+
{ type: 'agent_start' },
1076+
{ type: 'agent_end', messages: [makeUser('hi'), final], totalUsage: sampleUsage },
1077+
])
1078+
);
1079+
const { result } = renderHook(() => useChat({ api: '/chat', fetch: fetchFn }));
1080+
1081+
await act(async () => {
1082+
await result.current.send('hi');
1083+
});
1084+
1085+
expect(result.current.usage).not.toBeNull();
1086+
expect(result.current.usage?.input).toBe(10);
1087+
expect(result.current.usage?.output).toBe(20);
1088+
});
1089+
1090+
it('is populated from turn_end totalUsage', async () => {
1091+
const final = makeFinalAssistant('ok');
1092+
const fetchFn = jest.fn(
1093+
async (): Promise<Response> =>
1094+
streamFromEvents([
1095+
{ type: 'agent_start' },
1096+
{
1097+
type: 'turn_end',
1098+
message: final,
1099+
toolResults: [],
1100+
totalUsage: sampleUsage,
1101+
},
1102+
{ type: 'agent_end', messages: [makeUser('hi'), final], totalUsage: sampleUsage },
1103+
])
1104+
);
1105+
const { result } = renderHook(() => useChat({ api: '/chat', fetch: fetchFn }));
1106+
1107+
await act(async () => {
1108+
await result.current.send('hi');
1109+
});
1110+
1111+
expect(result.current.usage?.input).toBe(10);
1112+
expect(result.current.usage?.output).toBe(20);
1113+
});
1114+
1115+
it('resets to null when a new prompt() is called', async () => {
1116+
const final = makeFinalAssistant('ok');
1117+
let releaseFetch: (() => void) | null = null;
1118+
const fetchFn = jest.fn(
1119+
async (): Promise<Response> =>
1120+
streamFromEvents([
1121+
{ type: 'agent_start' },
1122+
{ type: 'agent_end', messages: [makeUser('hi'), final], totalUsage: sampleUsage },
1123+
])
1124+
);
1125+
1126+
// First, complete a request so usage is populated.
1127+
const { result } = renderHook(() => useChat({ api: '/chat', fetch: fetchFn }));
1128+
await act(async () => {
1129+
await result.current.send('first');
1130+
});
1131+
expect(result.current.usage?.input).toBe(10);
1132+
1133+
// Now set up a blocking fetch for the second request.
1134+
fetchFn.mockImplementationOnce(
1135+
async (): Promise<Response> => {
1136+
await new Promise<void>((resolve) => {
1137+
releaseFetch = resolve;
1138+
});
1139+
return streamFromEvents([
1140+
{ type: 'agent_start' },
1141+
{ type: 'agent_end', messages: [makeUser('second'), final], totalUsage: sampleUsage },
1142+
]);
1143+
}
1144+
);
1145+
1146+
// Start the second send (do not await yet).
1147+
act(() => {
1148+
void result.current.send('second');
1149+
});
1150+
1151+
// usage should be reset to null immediately after send() is called.
1152+
await waitFor(() => expect(result.current.usage).toBeNull());
1153+
1154+
// Release and finish.
1155+
await act(async () => {
1156+
releaseFetch?.();
1157+
});
1158+
});
1159+
});
1160+
10521161
describe('lifecycle hooks', () => {
10531162
it('reads handlers from a ref so consumers do not need to memoize', async () => {
10541163
const final = makeFinalAssistant('ok');
@@ -1060,7 +1169,7 @@ describe('useChat', () => {
10601169
{ type: 'message_end', message: makeUser('hi') },
10611170
{ type: 'message_start', message: final },
10621171
{ type: 'message_end', message: final },
1063-
{ type: 'agent_end', messages: [makeUser('hi'), final] },
1172+
{ type: 'agent_end', messages: [makeUser('hi'), final], totalUsage: ZERO_USAGE },
10641173
])
10651174
);
10661175

0 commit comments

Comments
 (0)