Skip to content

Commit bf87b33

Browse files
josers18claude
andcommitted
feat(usage): per-step token counts in reasoning trails (both stacks)
Reasoning trail rows now group under per-iteration headers showing the model's exact in/out tokens for that turn, with a ~N tok result-size estimate on each tool row and a total-tokens badge on the collapsed header. - main loop emits iteration_usage (per-iteration delta) + resultTokens on tool_result; ReasoningTrail groups + displays - second stack (ask-data + analyze) carries iteration on its usage event and resultTokens on tool_result; AskDataTrace groups + displays - shared forwardAgentEvent helper replaces the 10 identical inline route forwarders (also fixes error forwarding being main-stack-inconsistent) - review fixes: simplified dead gotExact expr; track last-defined iteration so an error-aggregate between same-turn rows can't emit a dup header Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent d57d5fa commit bf87b33

26 files changed

Lines changed: 530 additions & 312 deletions

File tree

app/api/analyze-ask/route.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,9 @@ export async function POST(req: NextRequest) {
144144
name: ev.name,
145145
isError: ev.isError,
146146
preview: ev.preview,
147+
...(typeof ev.resultTokens === "number"
148+
? { resultTokens: ev.resultTokens }
149+
: {}),
147150
});
148151
} else if (ev.type === "table_fallback") {
149152
send({
@@ -169,6 +172,13 @@ export async function POST(req: NextRequest) {
169172
usageAcc.inputTokens += ev.inputTokens;
170173
usageAcc.outputTokens += ev.outputTokens;
171174
usageAcc.exact = usageAcc.exact && ev.exact;
175+
send({
176+
type: "iteration_usage",
177+
iteration: ev.iteration,
178+
inputTokens: ev.inputTokens,
179+
outputTokens: ev.outputTokens,
180+
exact: ev.exact,
181+
});
172182
}
173183
}
174184
} finally {

app/api/arc/route.ts

Lines changed: 2 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { ensureFreshToken, getSessionId } from "@/lib/salesforce/token";
33
import { runAgentWithMcp } from "@/lib/llm/provider";
44
import { SYSTEM_PROMPT } from "@/lib/prompts/system";
55
import { arcPrompt } from "@/lib/prompts/arc";
6-
import { makeCacheableSseStream, sendInferenceMeta } from "@/lib/sse/stream";
6+
import { makeCacheableSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
77
import { localDayInTz } from "@/lib/sse/sectionCache";
88
import { log, correlationId } from "@/lib/log";
99
import { optionalEnv } from "@/lib/utils";
@@ -48,28 +48,7 @@ export async function GET(req: NextRequest) {
4848
maxTokens: 2048,
4949
forceFirstToolCall: true,
5050
routeHint: "arc",
51-
onEvent: (e) => {
52-
if (e.type === "text_delta" && e.text) {
53-
send({ type: "text_delta", text: e.text });
54-
} else if (e.type === "tool_use" && e.server && e.tool) {
55-
send({
56-
type: "tool_use",
57-
server: e.server,
58-
tool: e.tool,
59-
input: e.input,
60-
});
61-
} else if (e.type === "tool_result" && e.server && e.tool) {
62-
send({
63-
type: "tool_result",
64-
server: e.server,
65-
tool: e.tool,
66-
is_error: e.is_error,
67-
preview: e.preview ?? "",
68-
});
69-
} else if (e.type === "usage_meta" && e.usage) {
70-
send({ type: "usage_meta", usage: e.usage });
71-
}
72-
},
51+
onEvent: (e) => forwardAgentEvent(send, e),
7352
});
7453
sendInferenceMeta(send, result.inferenceBackend);
7554
log.info("arc.done", {

app/api/ask-data/route.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,9 @@ export async function POST(req: NextRequest) {
202202
name: ev.name,
203203
isError: ev.isError,
204204
preview: ev.preview,
205+
...(typeof ev.resultTokens === "number"
206+
? { resultTokens: ev.resultTokens }
207+
: {}),
205208
});
206209
} else if (ev.type === "turn_complete") {
207210
finalAssistantText = ev.text;
@@ -214,6 +217,14 @@ export async function POST(req: NextRequest) {
214217
usageAcc.inputTokens += ev.inputTokens;
215218
usageAcc.outputTokens += ev.outputTokens;
216219
usageAcc.exact = usageAcc.exact && ev.exact;
220+
// Forward per-iteration usage to the trail.
221+
send({
222+
type: "iteration_usage",
223+
iteration: ev.iteration,
224+
inputTokens: ev.inputTokens,
225+
outputTokens: ev.outputTokens,
226+
exact: ev.exact,
227+
});
217228
}
218229
}
219230
} finally {

app/api/ask/route.ts

Lines changed: 2 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { ensureFreshToken, getSessionId } from "@/lib/salesforce/token";
44
import { runAgentWithMcp } from "@/lib/llm/provider";
55
import { SYSTEM_PROMPT } from "@/lib/prompts/system";
66
import { askAnythingPrompt } from "@/lib/prompts/ask-anything";
7-
import { makeSseStream, sendInferenceMeta } from "@/lib/sse/stream";
7+
import { makeSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
88
import { log, correlationId } from "@/lib/log";
99
import { optionalEnv } from "@/lib/utils";
1010
import { validateAskThreadMessages } from "@/lib/ask/thread";
@@ -94,28 +94,7 @@ export async function POST(req: NextRequest) {
9494
maxIterations: 12,
9595
forceFirstToolCall,
9696
routeHint: fromGhostClick ? "ghost-ask" : undefined,
97-
onEvent: (e) => {
98-
if (e.type === "text_delta" && e.text) {
99-
send({ type: "text_delta", text: e.text });
100-
} else if (e.type === "tool_use" && e.server && e.tool) {
101-
send({
102-
type: "tool_use",
103-
server: e.server,
104-
tool: e.tool,
105-
input: e.input,
106-
});
107-
} else if (e.type === "tool_result" && e.server && e.tool) {
108-
send({
109-
type: "tool_result",
110-
server: e.server,
111-
tool: e.tool,
112-
is_error: e.is_error,
113-
preview: e.preview ?? "",
114-
});
115-
} else if (e.type === "usage_meta" && e.usage) {
116-
send({ type: "usage_meta", usage: e.usage });
117-
}
118-
},
97+
onEvent: (e) => forwardAgentEvent(send, e),
11998
});
12099
sendInferenceMeta(send, result.inferenceBackend);
121100
send({ type: "thread_snapshot", messages: result.transcript });

app/api/brief/route.ts

Lines changed: 2 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { ensureFreshToken, getSessionId, resolveBankerDisplayName } from "@/lib/
33
import { runAgentWithMcp } from "@/lib/llm/provider";
44
import { SYSTEM_PROMPT } from "@/lib/prompts/system";
55
import { morningBriefPrompt } from "@/lib/prompts/morning-brief";
6-
import { makeCacheableSseStream, sendInferenceMeta } from "@/lib/sse/stream";
6+
import { makeCacheableSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
77
import { localDayInTz } from "@/lib/sse/sectionCache";
88
import { log, correlationId } from "@/lib/log";
99
import { optionalEnv } from "@/lib/utils";
@@ -71,34 +71,7 @@ export async function POST(req: NextRequest) {
7171
maxIterations: 6,
7272
maxTokens: 2048,
7373
routeHint: "morning-brief",
74-
onEvent: (e) => {
75-
if (e.type === "text_delta" && e.text) {
76-
send({ type: "text_delta", text: e.text });
77-
} else if (
78-
e.type === "error" &&
79-
typeof e.message === "string" &&
80-
e.message.length > 0
81-
) {
82-
send({ type: "error", message: e.message });
83-
} else if (e.type === "tool_use" && e.server && e.tool) {
84-
send({
85-
type: "tool_use",
86-
server: e.server,
87-
tool: e.tool,
88-
input: e.input,
89-
});
90-
} else if (e.type === "tool_result" && e.server && e.tool) {
91-
send({
92-
type: "tool_result",
93-
server: e.server,
94-
tool: e.tool,
95-
is_error: e.is_error,
96-
preview: e.preview ?? "",
97-
});
98-
} else if (e.type === "usage_meta" && e.usage) {
99-
send({ type: "usage_meta", usage: e.usage });
100-
}
101-
},
74+
onEvent: (e) => forwardAgentEvent(send, e),
10275
});
10376
sendInferenceMeta(send, result.inferenceBackend);
10477
log.info("brief.done", {

app/api/client/[id]/route.ts

Lines changed: 2 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { ensureFreshToken, getSessionId } from "@/lib/salesforce/token";
33
import { runAgentWithMcp } from "@/lib/llm/provider";
44
import { SYSTEM_PROMPT } from "@/lib/prompts/system";
55
import { clientDetailPrompt } from "@/lib/prompts/client-detail";
6-
import { makeSseStream, sendInferenceMeta } from "@/lib/sse/stream";
6+
import { makeSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
77
import { log, correlationId } from "@/lib/log";
88
import { optionalEnv } from "@/lib/utils";
99

@@ -56,28 +56,7 @@ export async function GET(
5656
maxIterations: 4,
5757
forceFirstToolCall: true,
5858
routeHint: "client-detail",
59-
onEvent: (e) => {
60-
if (e.type === "text_delta" && e.text) {
61-
send({ type: "text_delta", text: e.text });
62-
} else if (e.type === "tool_use" && e.server && e.tool) {
63-
send({
64-
type: "tool_use",
65-
server: e.server,
66-
tool: e.tool,
67-
input: e.input,
68-
});
69-
} else if (e.type === "tool_result" && e.server && e.tool) {
70-
send({
71-
type: "tool_result",
72-
server: e.server,
73-
tool: e.tool,
74-
is_error: e.is_error,
75-
preview: e.preview ?? "",
76-
});
77-
} else if (e.type === "usage_meta" && e.usage) {
78-
send({ type: "usage_meta", usage: e.usage });
79-
}
80-
},
59+
onEvent: (e) => forwardAgentEvent(send, e),
8160
});
8261
sendInferenceMeta(send, result.inferenceBackend);
8362
log.info("client.done", {

app/api/drafts/route.ts

Lines changed: 2 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { ensureFreshToken, getSessionId } from "@/lib/salesforce/token";
33
import { runAgentWithMcp } from "@/lib/llm/provider";
44
import { SYSTEM_PROMPT } from "@/lib/prompts/system";
55
import { draftQueuePrompt } from "@/lib/prompts/draft-queue";
6-
import { makeCacheableSseStream, sendInferenceMeta } from "@/lib/sse/stream";
6+
import { makeCacheableSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
77
import { localDayInTz } from "@/lib/sse/sectionCache";
88
import { log, correlationId } from "@/lib/log";
99
import { optionalEnv } from "@/lib/utils";
@@ -58,28 +58,7 @@ export async function GET(req: NextRequest) {
5858
maxIterations: 7,
5959
maxTokens: 3072,
6060
routeHint: "drafts",
61-
onEvent: (e) => {
62-
if (e.type === "text_delta" && e.text) {
63-
send({ type: "text_delta", text: e.text });
64-
} else if (e.type === "tool_use" && e.server && e.tool) {
65-
send({
66-
type: "tool_use",
67-
server: e.server,
68-
tool: e.tool,
69-
input: e.input,
70-
});
71-
} else if (e.type === "tool_result" && e.server && e.tool) {
72-
send({
73-
type: "tool_result",
74-
server: e.server,
75-
tool: e.tool,
76-
is_error: e.is_error,
77-
preview: e.preview ?? "",
78-
});
79-
} else if (e.type === "usage_meta" && e.usage) {
80-
send({ type: "usage_meta", usage: e.usage });
81-
}
82-
},
61+
onEvent: (e) => forwardAgentEvent(send, e),
8362
});
8463
sendInferenceMeta(send, result.inferenceBackend);
8564
log.info("drafts.done", {

app/api/insights/route.ts

Lines changed: 3 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import {
77
sectionInsightPrompt,
88
type SectionKind,
99
} from "@/lib/prompts/section-insight";
10-
import { makeSseStream, sendInferenceMeta } from "@/lib/sse/stream";
10+
import { makeSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
1111
import { log, correlationId } from "@/lib/log";
1212
import { optionalEnv } from "@/lib/utils";
1313

@@ -70,28 +70,7 @@ export async function POST(req: NextRequest) {
7070
maxIterations: 6,
7171
maxTokens: 2048,
7272
routeHint: "insights",
73-
onEvent: (e) => {
74-
if (e.type === "text_delta" && e.text) {
75-
send({ type: "text_delta", text: e.text });
76-
} else if (e.type === "tool_use" && e.server && e.tool) {
77-
send({
78-
type: "tool_use",
79-
server: e.server,
80-
tool: e.tool,
81-
input: e.input,
82-
});
83-
} else if (e.type === "tool_result" && e.server && e.tool) {
84-
send({
85-
type: "tool_result",
86-
server: e.server,
87-
tool: e.tool,
88-
is_error: e.is_error,
89-
preview: e.preview ?? "",
90-
});
91-
} else if (e.type === "usage_meta" && e.usage) {
92-
send({ type: "usage_meta", usage: e.usage });
93-
}
94-
},
73+
onEvent: (e) => forwardAgentEvent(send, e),
9574
});
9675
sendInferenceMeta(send, result.inferenceBackend);
9776
log.info("insight.batch.done", {
@@ -126,28 +105,7 @@ export async function POST(req: NextRequest) {
126105
maxIterations: 4,
127106
maxTokens: 1024,
128107
routeHint: "insights",
129-
onEvent: (e) => {
130-
if (e.type === "text_delta" && e.text) {
131-
send({ type: "text_delta", text: e.text });
132-
} else if (e.type === "tool_use" && e.server && e.tool) {
133-
send({
134-
type: "tool_use",
135-
server: e.server,
136-
tool: e.tool,
137-
input: e.input,
138-
});
139-
} else if (e.type === "tool_result" && e.server && e.tool) {
140-
send({
141-
type: "tool_result",
142-
server: e.server,
143-
tool: e.tool,
144-
is_error: e.is_error,
145-
preview: e.preview ?? "",
146-
});
147-
} else if (e.type === "usage_meta" && e.usage) {
148-
send({ type: "usage_meta", usage: e.usage });
149-
}
150-
},
108+
onEvent: (e) => forwardAgentEvent(send, e),
151109
});
152110
sendInferenceMeta(send, result.inferenceBackend);
153111
log.info("insight.done", {

app/api/prep/route.ts

Lines changed: 2 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { ensureFreshToken, getSessionId, resolveBankerDisplayName } from "@/lib/
33
import { runAgentWithMcp } from "@/lib/llm/provider";
44
import { SYSTEM_PROMPT } from "@/lib/prompts/system";
55
import { prepPrompt } from "@/lib/prompts/prep";
6-
import { makeSseStream, sendInferenceMeta } from "@/lib/sse/stream";
6+
import { makeSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
77
import { log, correlationId } from "@/lib/log";
88
import { optionalEnv } from "@/lib/utils";
99

@@ -59,28 +59,7 @@ export async function POST(req: NextRequest) {
5959
route: "prep",
6060
maxIterations: 14,
6161
routeHint: "prep",
62-
onEvent: (e) => {
63-
if (e.type === "text_delta" && e.text) {
64-
send({ type: "text_delta", text: e.text });
65-
} else if (e.type === "tool_use" && e.server && e.tool) {
66-
send({
67-
type: "tool_use",
68-
server: e.server,
69-
tool: e.tool,
70-
input: e.input,
71-
});
72-
} else if (e.type === "tool_result" && e.server && e.tool) {
73-
send({
74-
type: "tool_result",
75-
server: e.server,
76-
tool: e.tool,
77-
is_error: e.is_error,
78-
preview: e.preview ?? "",
79-
});
80-
} else if (e.type === "usage_meta" && e.usage) {
81-
send({ type: "usage_meta", usage: e.usage });
82-
}
83-
},
62+
onEvent: (e) => forwardAgentEvent(send, e),
8463
});
8564
sendInferenceMeta(send, result.inferenceBackend);
8665
log.info("prep.done", {

app/api/priority/route.ts

Lines changed: 2 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { ensureFreshToken, getSessionId } from "@/lib/salesforce/token";
33
import { runAgentWithMcp } from "@/lib/llm/provider";
44
import { SYSTEM_PROMPT } from "@/lib/prompts/system";
55
import { priorityQueuePrompt } from "@/lib/prompts/priority-queue";
6-
import { makeCacheableSseStream, sendInferenceMeta } from "@/lib/sse/stream";
6+
import { makeCacheableSseStream, sendInferenceMeta, forwardAgentEvent } from "@/lib/sse/stream";
77
import { localDayInTz } from "@/lib/sse/sectionCache";
88
import { log, correlationId } from "@/lib/log";
99
import { optionalEnv } from "@/lib/utils";
@@ -62,28 +62,7 @@ export async function GET(req: NextRequest) {
6262
maxIterations: 7,
6363
maxTokens: 2048,
6464
routeHint: "priority",
65-
onEvent: (e) => {
66-
if (e.type === "text_delta" && e.text) {
67-
send({ type: "text_delta", text: e.text });
68-
} else if (e.type === "tool_use" && e.server && e.tool) {
69-
send({
70-
type: "tool_use",
71-
server: e.server,
72-
tool: e.tool,
73-
input: e.input,
74-
});
75-
} else if (e.type === "tool_result" && e.server && e.tool) {
76-
send({
77-
type: "tool_result",
78-
server: e.server,
79-
tool: e.tool,
80-
is_error: e.is_error,
81-
preview: e.preview ?? "",
82-
});
83-
} else if (e.type === "usage_meta" && e.usage) {
84-
send({ type: "usage_meta", usage: e.usage });
85-
}
86-
},
65+
onEvent: (e) => forwardAgentEvent(send, e),
8766
});
8867
sendInferenceMeta(send, result.inferenceBackend);
8968
log.info("priority.done", {

0 commit comments

Comments
 (0)