feat: token economics — real API token extraction, stats propagation, and context window fuel gauge

Phase 1: OpenAI API Token Extraction
- Add stream_options: { include_usage: true } to all streaming API calls
- Capture chunk.usage from final streaming chunks and response.usage from non-streaming
- Extend OpenAiChatIterationResult with optional usage field
- Update buildStats() to accept and return real token counts from usage data
- Wire iteration.usage through chat() to both buildStats() call sites

Phase 2: Agent Loop Stats Propagation
- Rename inputTokens/outputTokens to masterInputTokens/masterOutputTokens, add masterThinkingTokens
- Accumulate response.stats.tokenCounts after each master AI call
- Delete all Math.ceil(length/4) crude approximations (master and subagent loops)
- Track startTime/durationMs and emit IWorkOrderCompleteStats with workOrderComplete
- Subagent loop uses response.stats?.tokenCounts instead of Math.ceil

Phase 3: Database Model Changes
- Add contextWindowUsage field to IChatSession, ChatSessionSchema, and frontend ChatSession
- Initialize contextWindowUsage: 0 on session creation

Phase 4: Persist Stats on Turn Completion
- drone-session: accept IWorkOrderCompleteStats, persist turn stats, walk subagent records for aggregate
- drone-session: $inc session stats and contextWindowUsage, add formatDurationLabel() helper
- code-session: accept and forward stats, update in-memory session stats
- message-queue: Redis replay handles 4th stats arg
- Update WorkOrderCompleteMessage type in @gadget/api to accept stats parameter

Phase 5: UI — Context Window Fuel Gauge
- Add contextWindowUsage prop to SessionPanel
- Add fuel gauge bar with E→F labels, green/yellow/red zones, token count display
- Visible for ALL provider types (not gated by apiType)

Phase 6: Frontend Streaming State
- Add IWorkOrderCompleteStats interface to frontend api.ts
- handleWorkOrderComplete accepts stats, updates turn stats and session contextWindowUsage
- Pass contextWindowUsage prop to SessionPanel
This commit is contained in:
Rob Colbert 2026-05-18 14:36:36 -04:00
parent 9c1e65785d
commit 00969cf9bc
12 changed files with 239 additions and 24 deletions

View File

@ -41,6 +41,7 @@ interface SessionPanelProps {
truncateModelName: (name: string, maxLength?: number) => string; truncateModelName: (name: string, maxLength?: number) => string;
getSelectedModelCapabilities: () => { hasThinking: boolean } | null; getSelectedModelCapabilities: () => { hasThinking: boolean } | null;
getSelectedModel: () => AiModel | null; getSelectedModel: () => AiModel | null;
contextWindowUsage: number;
} }
export default function SessionPanel({ export default function SessionPanel({
@ -82,7 +83,9 @@ export default function SessionPanel({
truncateModelName, truncateModelName,
getSelectedModelCapabilities, getSelectedModelCapabilities,
getSelectedModel, getSelectedModel,
contextWindowUsage,
}: SessionPanelProps) { }: SessionPanelProps) {
const usagePercent = sessionNumCtx > 0 ? (contextWindowUsage / sessionNumCtx) * 100 : 0;
return ( return (
<div className="border-b border-border-subtle shrink-0"> <div className="border-b border-border-subtle shrink-0">
<div className="flex items-center justify-between px-4 py-2 bg-bg-tertiary"> <div className="flex items-center justify-between px-4 py-2 bg-bg-tertiary">
@ -292,6 +295,33 @@ export default function SessionPanel({
</div> </div>
</div> </div>
)} )}
<div>
<div className="text-xs text-text-muted">Context Window Usage</div>
<div className="mt-1">
<div className="flex items-center gap-1.5">
<span className="text-[10px] font-mono text-text-muted w-3">E</span>
<div className="flex-1 h-3 bg-bg-tertiary rounded-full overflow-hidden border border-border-subtle">
<div
className={`h-full rounded-full transition-all duration-300 ${
usagePercent >= 80 ? 'bg-red-500' :
usagePercent >= 50 ? 'bg-yellow-500' :
'bg-green-500'
}`}
style={{ width: `${Math.min(usagePercent, 100)}%` }}
/>
</div>
<span className="text-[10px] font-mono text-text-muted w-3">F</span>
</div>
<div className="flex justify-between mt-0.5">
<span className="text-[10px] text-text-muted tabular-nums">
{contextWindowUsage.toLocaleString()} tokens
</span>
<span className="text-[10px] text-text-muted tabular-nums">
{sessionNumCtx.toLocaleString()} max
</span>
</div>
</div>
</div>
</div> </div>
</div> </div>
); );

View File

@ -388,6 +388,7 @@ export interface ChatSession {
selectedModel: string; selectedModel: string;
reasoningEffort?: string; reasoningEffort?: string;
numCtx?: number; numCtx?: number;
contextWindowUsage: number;
stats: ChatSessionStats; stats: ChatSessionStats;
pins: Array<{ _id?: string; content: string }>; pins: Array<{ _id?: string; content: string }>;
} }
@ -401,6 +402,14 @@ export interface ChatTurnStats {
durationLabel: string; durationLabel: string;
} }
export interface IWorkOrderCompleteStats {
masterInputTokens: number;
masterOutputTokens: number;
masterThinkingTokens: number;
toolCallCount: number;
durationMs: number;
}
export interface ChatTurnPrompts { export interface ChatTurnPrompts {
user: string; user: string;
system?: string; system?: string;

View File

@ -2,7 +2,7 @@ import { useState, useEffect, useRef, useContext, useCallback } from 'react';
import { useParams, useNavigate, useLocation } from 'react-router-dom'; import { useParams, useNavigate, useLocation } from 'react-router-dom';
import { BrushCleaning, Maximize2 } from 'lucide-react'; import { BrushCleaning, Maximize2 } from 'lucide-react';
import { socketClient } from '../lib/socket'; import { socketClient } from '../lib/socket';
import { chatSessionApi, projectApi, providerApi, type ChatSession, type ChatTurn, type ChatTurnBlock, ChatTurnStats, ChatSessionMode, type AiProvider, type Project } from '../lib/api'; import { chatSessionApi, projectApi, providerApi, type ChatSession, type ChatTurn, type ChatTurnBlock, ChatTurnStats, ChatSessionMode, type AiProvider, type Project, type IWorkOrderCompleteStats } from '../lib/api';
import { WorkspaceMode } from '../lib/types'; import { WorkspaceMode } from '../lib/types';
import SessionPanel from '../components/SessionPanel'; import SessionPanel from '../components/SessionPanel';
import ProjectPanel from '../components/ProjectPanel'; import ProjectPanel from '../components/ProjectPanel';
@ -686,7 +686,7 @@ export default function ChatSessionView() {
}, 4000); }, 4000);
}, []); }, []);
const handleWorkOrderComplete = useCallback((turnId: string, success: boolean, message?: string) => { const handleWorkOrderComplete = useCallback((turnId: string, success: boolean, message?: string, stats?: IWorkOrderCompleteStats) => {
// Backend has already flushed and persisted all streaming content // Backend has already flushed and persisted all streaming content
// Just clean up frontend streaming state and update status // Just clean up frontend streaming state and update status
if (streamingStateRef.current.has(turnId)) { if (streamingStateRef.current.has(turnId)) {
@ -734,6 +734,38 @@ export default function ChatSessionView() {
setIsAborting(false); setIsAborting(false);
setIsProcessing(false); setIsProcessing(false);
currentTurnIdRef.current = null; currentTurnIdRef.current = null;
// Update turn stats from completion event
if (stats) {
setTurns(prev => prev.map(t => {
if (t._id !== turnId) return t;
return {
...t,
stats: {
...t.stats,
inputTokens: stats.masterInputTokens,
thinkingTokenCount: stats.masterThinkingTokens,
responseTokens: stats.masterOutputTokens,
toolCallCount: stats.toolCallCount,
durationMs: stats.durationMs,
durationLabel: stats.durationMs < 1000
? `${stats.durationMs}ms`
: `${(stats.durationMs / 1000).toFixed(1)}s`,
},
};
}));
// Update session contextWindowUsage for fuel gauge
setSession(prev => prev ? {
...prev,
contextWindowUsage: (prev.contextWindowUsage ?? 0) + stats.masterInputTokens,
stats: {
...prev.stats,
inputTokens: prev.stats.inputTokens + stats.masterInputTokens,
outputTokens: prev.stats.outputTokens + stats.masterOutputTokens,
},
} : prev);
}
}, [showToast]); }, [showToast]);
const handleWorkspaceModeChanged = useCallback((mode: string) => { const handleWorkspaceModeChanged = useCallback((mode: string) => {
@ -1441,6 +1473,7 @@ export default function ChatSessionView() {
truncateModelName={truncateModelName} truncateModelName={truncateModelName}
getSelectedModelCapabilities={getSelectedModelCapabilities} getSelectedModelCapabilities={getSelectedModelCapabilities}
getSelectedModel={getSelectedModel} getSelectedModel={getSelectedModel}
contextWindowUsage={session?.contextWindowUsage ?? 0}
/> />
<ProjectPanel <ProjectPanel

View File

@ -23,6 +23,7 @@ import {
SubmitPromptCallback, SubmitPromptCallback,
AbortWorkOrderCallback, AbortWorkOrderCallback,
SubProcessStat, SubProcessStat,
IWorkOrderCompleteStats,
} from "@gadget/api"; } from "@gadget/api";
import ChatSession from "../models/chat-session.ts"; import ChatSession from "../models/chat-session.ts";
@ -669,7 +670,16 @@ export class CodeSession extends SocketSession {
turnId: string, turnId: string,
success: boolean, success: boolean,
message?: string, message?: string,
stats?: IWorkOrderCompleteStats,
): void { ): void {
this.socket.emit("workOrderComplete", turnId, success, message); this.socket.emit("workOrderComplete", turnId, success, message, stats);
// Update in-memory session stats for UI reactivity
if (stats && this.chatSession) {
this.chatSession.stats.inputTokens += stats.masterInputTokens;
this.chatSession.stats.outputTokens += stats.masterOutputTokens;
this.chatSession.stats.toolCallCount += stats.toolCallCount;
this.chatSession.contextWindowUsage = (this.chatSession.contextWindowUsage ?? 0) + stats.masterInputTokens;
}
} }
} }

View File

@ -12,6 +12,7 @@ import {
WorkspaceMode, WorkspaceMode,
IChatToolCall, IChatToolCall,
IChatTurnBlock, IChatTurnBlock,
IWorkOrderCompleteStats,
} from "@gadget/api"; } from "@gadget/api";
import { import {
GadgetSocket, GadgetSocket,
@ -20,6 +21,7 @@ import {
} from "./socket-session.js"; } from "./socket-session.js";
import { SocketService } from "../services/index.js"; import { SocketService } from "../services/index.js";
import { ChatTurn } from "../models/chat-turn.js"; import { ChatTurn } from "../models/chat-turn.js";
import ChatSession from "../models/chat-session.js";
import MessageQueue, { type QueuedMessage } from "./message-queue.js"; import MessageQueue, { type QueuedMessage } from "./message-queue.js";
interface IStreamingBuffer { interface IStreamingBuffer {
@ -292,6 +294,7 @@ export class DroneSession extends SocketSession {
turnId: string, turnId: string,
success: boolean, success: boolean,
message?: string, message?: string,
stats?: IWorkOrderCompleteStats,
): Promise<void> { ): Promise<void> {
if (!this.chatSessionId) { if (!this.chatSessionId) {
this.log.warn("workOrderComplete event received but no chat session is active"); this.log.warn("workOrderComplete event received but no chat session is active");
@ -314,20 +317,55 @@ export class DroneSession extends SocketSession {
turn.errorMessage = message; turn.errorMessage = message;
} }
} }
// UPDATE TURN STATS — master agent tokens only
if (stats) {
turn.stats.inputTokens = stats.masterInputTokens;
turn.stats.thinkingTokenCount = stats.masterThinkingTokens;
turn.stats.responseTokens = stats.masterOutputTokens;
turn.stats.toolCallCount = stats.toolCallCount;
turn.stats.durationMs = stats.durationMs;
turn.stats.durationLabel = this.formatDurationLabel(stats.durationMs);
}
await turn.save(); await turn.save();
// COMPUTE AGGREGATE: walk the turn's subagent records for aggregate totals
let totalInputTokens = stats?.masterInputTokens ?? 0;
let totalOutputTokens = stats?.masterOutputTokens ?? 0;
if (stats) {
for (const tc of turn.toolCalls) {
if ((tc as any).subagent?.stats) {
totalInputTokens += (tc as any).subagent.stats.inputTokens;
totalOutputTokens += (tc as any).subagent.stats.responseTokens + (tc as any).subagent.stats.thinkingTokenCount;
}
}
}
// INCREMENT SESSION STATS — aggregate totals (master + subagents)
await ChatSession.updateOne(
{ _id: this.chatSessionId },
{ $inc: {
"stats.inputTokens": totalInputTokens,
"stats.outputTokens": totalOutputTokens,
"stats.toolCallCount": stats?.toolCallCount ?? 0,
// Master-only field for the fuel gauge
"contextWindowUsage": stats?.masterInputTokens ?? 0,
}},
);
} }
const codeSession = SocketService.getCodeSessionByChatSessionId( const codeSession = SocketService.getCodeSessionByChatSessionId(
this.chatSessionId, this.chatSessionId,
); );
codeSession.onWorkOrderComplete(turnId, success, message); codeSession.onWorkOrderComplete(turnId, success, message, stats);
this.currentTurnId = undefined; this.currentTurnId = undefined;
} catch (error) { } catch (error) {
// Routing failed - queue to Redis // Routing failed - queue to Redis
await MessageQueue.enqueue(this.chatSessionId, { await MessageQueue.enqueue(this.chatSessionId, {
type: 'workOrderComplete', type: 'workOrderComplete',
args: [turnId, success, message], args: [turnId, success, message, stats],
timestamp: Date.now(), timestamp: Date.now(),
}); });
this.log.debug("queued workOrderComplete message", { chatSessionId: this.chatSessionId }); this.log.debug("queued workOrderComplete message", { chatSessionId: this.chatSessionId });
@ -343,6 +381,17 @@ export class DroneSession extends SocketSession {
this.streamingBuffers.clear(); this.streamingBuffers.clear();
} }
private formatDurationLabel(durationMs: number): string {
if (durationMs < 1000) return `${durationMs}ms`;
const seconds = Math.floor(durationMs / 1000);
const minutes = Math.floor(seconds / 60);
const secs = seconds % 60;
if (minutes > 0) {
return `${minutes}m ${secs}s`;
}
return `${seconds}.${Math.floor((durationMs % 1000) / 100)}s`;
}
/** /**
* Gets or creates a streaming buffer for a turn. * Gets or creates a streaming buffer for a turn.
*/ */
@ -669,7 +718,8 @@ export class DroneSession extends SocketSession {
codeSession.onWorkOrderComplete( codeSession.onWorkOrderComplete(
msg.args[0] as string, msg.args[0] as string,
msg.args[1] as boolean, msg.args[1] as boolean,
msg.args[2] as string, msg.args[2] as string | undefined,
msg.args[3] as IWorkOrderCompleteStats | undefined,
); );
break; break;
case 'log': case 'log':

View File

@ -31,6 +31,7 @@ export const ChatSessionSchema = new Schema<IChatSession>({
default: "off", default: "off",
}, },
numCtx: { type: Number }, numCtx: { type: Number },
contextWindowUsage: { type: Number, default: 0, required: true },
stats: { stats: {
turnCount: { type: Number, default: 0, required: true }, turnCount: { type: Number, default: 0, required: true },
toolCallCount: { type: Number, default: 0, required: true }, toolCallCount: { type: Number, default: 0, required: true },

View File

@ -131,6 +131,7 @@ class ChatSessionService extends DtpService {
inputTokens: 0, inputTokens: 0,
outputTokens: 0, outputTokens: 0,
}, },
contextWindowUsage: 0,
pins: [], pins: [],
}); });

View File

@ -39,6 +39,7 @@ describe("AgentService", () => {
inputTokens: 0, inputTokens: 0,
outputTokens: 0, outputTokens: 0,
}, },
contextWindowUsage: 0,
pins: [], pins: [],
}; };
const workOrder: IAgentWorkOrder = { const workOrder: IAgentWorkOrder = {
@ -150,6 +151,7 @@ describe("AgentService", () => {
inputTokens: 0, inputTokens: 0,
outputTokens: 0, outputTokens: 0,
}, },
contextWindowUsage: 0,
pins: [], pins: [],
}; };
@ -314,6 +316,7 @@ describe("AgentService", () => {
inputTokens: 0, inputTokens: 0,
outputTokens: 0, outputTokens: 0,
}, },
contextWindowUsage: 0,
pins: [], pins: [],
}; };
const rawResponse = "Subagent completed the task successfully."; const rawResponse = "Subagent completed the task successfully.";

View File

@ -58,6 +58,14 @@ import {
ReadSkillTool, ReadSkillTool,
type DroneToolboxEnvironment, type DroneToolboxEnvironment,
} from "../tools/index.ts"; } from "../tools/index.ts";
interface IWorkOrderCompleteStats {
masterInputTokens: number;
masterOutputTokens: number;
masterThinkingTokens: number;
toolCallCount: number;
durationMs: number;
}
import type { IChatExportData } from "@gadget/ai-toolbox"; import type { IChatExportData } from "@gadget/ai-toolbox";
export interface IAgentWorkOrder { export interface IAgentWorkOrder {
@ -164,8 +172,10 @@ class AgentService extends GadgetService {
const { turn } = workOrder; const { turn } = workOrder;
let toolCallCount = 0; let toolCallCount = 0;
let inputTokens = 0; let masterInputTokens = 0;
let outputTokens = 0; let masterOutputTokens = 0;
let masterThinkingTokens = 0;
const startTime = Date.now();
// Build the full message array that grows with each iteration // Build the full message array that grows with each iteration
const messages: IContextChatMessage[] = []; const messages: IContextChatMessage[] = [];
@ -258,6 +268,11 @@ class AgentService extends GadgetService {
throw new Error("Model failed to respond (still loading or error)"); throw new Error("Model failed to respond (still loading or error)");
} }
// Accumulate real token counts from the AI API response
masterInputTokens += response.stats.tokenCounts.input;
masterOutputTokens += response.stats.tokenCounts.response;
masterThinkingTokens += response.stats.tokenCounts.thinking;
// Process tool calls if present // Process tool calls if present
if (response.toolCalls && response.toolCalls.length > 0) { if (response.toolCalls && response.toolCalls.length > 0) {
continueLoop = true; continueLoop = true;
@ -338,17 +353,30 @@ class AgentService extends GadgetService {
toolName: toolCall.function.name, toolName: toolCall.function.name,
}); });
inputTokens += Math.ceil(toolCall.function.arguments.length / 4);
outputTokens += Math.ceil(result.length / 4);
} }
} }
} }
socket.emit("workOrderComplete", turn._id, true); const durationMs = Date.now() - startTime;
const stats: IWorkOrderCompleteStats = {
masterInputTokens,
masterOutputTokens,
masterThinkingTokens,
toolCallCount,
durationMs,
};
socket.emit("workOrderComplete", turn._id, true, undefined, stats);
} catch (cause) { } catch (cause) {
if (cause instanceof Error && cause.name === "AbortError") { if (cause instanceof Error && cause.name === "AbortError") {
this.log.info("work order aborted by user", { turnId: turn._id }); this.log.info("work order aborted by user", { turnId: turn._id });
socket.emit("workOrderComplete", turn._id, true, "aborted"); const abortStats: IWorkOrderCompleteStats = {
masterInputTokens,
masterOutputTokens,
masterThinkingTokens,
toolCallCount,
durationMs: Date.now() - startTime,
};
socket.emit("workOrderComplete", turn._id, true, "aborted", abortStats);
return; return;
} }
const msg = cause instanceof Error ? cause.message : String(cause); const msg = cause instanceof Error ? cause.message : String(cause);
@ -356,7 +384,14 @@ class AgentService extends GadgetService {
turnId: turn._id, turnId: turn._id,
error: msg, error: msg,
}); });
socket.emit("workOrderComplete", turn._id, false, msg); const errorStats: IWorkOrderCompleteStats = {
masterInputTokens,
masterOutputTokens,
masterThinkingTokens,
toolCallCount,
durationMs: Date.now() - startTime,
};
socket.emit("workOrderComplete", turn._id, false, msg, errorStats);
throw cause; throw cause;
} finally { } finally {
this.abortController = null; this.abortController = null;
@ -740,11 +775,11 @@ class AgentService extends GadgetService {
inputTokens += response.stats?.tokenCounts?.input ?? 0; inputTokens += response.stats?.tokenCounts?.input ?? 0;
if (response.thinking) { if (response.thinking) {
thinkingTokenCount += Math.ceil(response.thinking.length / 4); thinkingTokenCount += response.stats?.tokenCounts?.thinking ?? 0;
} }
if (response.response) { if (response.response) {
fullResponse += response.response; fullResponse += response.response;
outputTokens += Math.ceil(response.response.length / 4); outputTokens += response.stats?.tokenCounts?.response ?? 0;
} }
if (response.thinking) { if (response.thinking) {
fullThinking = (fullThinking + response.thinking).trim(); fullThinking = (fullThinking + response.thinking).trim();
@ -800,8 +835,6 @@ class AgentService extends GadgetService {
toolName: toolCall.function.name, toolName: toolCall.function.name,
}); });
inputTokens += Math.ceil(toolArgsRaw.length / 4);
outputTokens += Math.ceil(result.length / 4);
} }
} }
} }

View File

@ -80,6 +80,11 @@ interface OpenAiChatIterationResult {
chunkCount: number; chunkCount: number;
contentDeltaCount: number; contentDeltaCount: number;
toolDeltaCount: number; toolDeltaCount: number;
// Usage data from OpenAI API (prompt_tokens, completion_tokens)
usage?: {
promptTokens: number;
completionTokens: number;
};
} }
export class OpenAiApi extends AiApi { export class OpenAiApi extends AiApi {
@ -206,6 +211,7 @@ export class OpenAiApi extends AiApi {
{ role: "user" as const, content: options.prompt }, { role: "user" as const, content: options.prompt },
], ],
stream: true, stream: true,
stream_options: { include_usage: true },
...(model.params.maxCompletionTokens ...(model.params.maxCompletionTokens
? { max_completion_tokens: model.params.maxCompletionTokens } ? { max_completion_tokens: model.params.maxCompletionTokens }
: {}), : {}),
@ -223,12 +229,21 @@ export class OpenAiApi extends AiApi {
let accumulatedResponse = ""; let accumulatedResponse = "";
let accumulatedThinking = ""; let accumulatedThinking = "";
let usage: { promptTokens: number; completionTokens: number } | undefined;
for await (const chunk of response) { for await (const chunk of response) {
if (options.signal?.aborted) { if (options.signal?.aborted) {
throw new DOMException("The operation was aborted", "AbortError"); throw new DOMException("The operation was aborted", "AbortError");
} }
// Capture usage from the final streaming chunk
if (chunk.usage) {
usage = {
promptTokens: chunk.usage.prompt_tokens,
completionTokens: chunk.usage.completion_tokens,
};
}
const delta = chunk.choices[0]?.delta; const delta = chunk.choices[0]?.delta;
if (delta) { if (delta) {
if (delta.content) { if (delta.content) {
@ -265,8 +280,8 @@ export class OpenAiApi extends AiApi {
text: numeral(durationMs / 1000).format("hh:mm:ss"), text: numeral(durationMs / 1000).format("hh:mm:ss"),
}, },
tokenCounts: { tokenCounts: {
input: 0, input: usage?.promptTokens ?? 0,
response: 0, response: usage?.completionTokens ?? 0,
thinking: 0, thinking: 0,
}, },
}, },
@ -325,7 +340,7 @@ export class OpenAiApi extends AiApi {
thinking: undefined, thinking: undefined,
toolCalls: undefined, toolCalls: undefined,
toolCallResults: undefined, toolCallResults: undefined,
stats: this.buildStats(startTime), stats: this.buildStats(startTime, iteration.usage),
}); });
} }
@ -334,7 +349,7 @@ export class OpenAiApi extends AiApi {
response: iteration.response, response: iteration.response,
thinking: iteration.thinking, thinking: iteration.thinking,
toolCalls: iteration.toolCalls.length > 0 ? iteration.toolCalls : undefined, toolCalls: iteration.toolCalls.length > 0 ? iteration.toolCalls : undefined,
stats: this.buildStats(startTime), stats: this.buildStats(startTime, iteration.usage),
}; };
this.assertNonEmptyChatResponse(finalResponse); this.assertNonEmptyChatResponse(finalResponse);
return finalResponse; return finalResponse;
@ -401,6 +416,7 @@ export class OpenAiApi extends AiApi {
messages, messages,
tools, tools,
stream: true, stream: true,
stream_options: { include_usage: true },
...(model.params.maxCompletionTokens ...(model.params.maxCompletionTokens
? { max_completion_tokens: model.params.maxCompletionTokens } ? { max_completion_tokens: model.params.maxCompletionTokens }
: {}), : {}),
@ -423,6 +439,7 @@ export class OpenAiApi extends AiApi {
let toolDeltaCount = 0; let toolDeltaCount = 0;
let finishReason: string | null | undefined; let finishReason: string | null | undefined;
const toolCallMap = new Map<number, StreamingToolCallAccumulator>(); const toolCallMap = new Map<number, StreamingToolCallAccumulator>();
let usage: { promptTokens: number; completionTokens: number } | undefined;
for await (const chunk of response) { for await (const chunk of response) {
if (signal?.aborted) { if (signal?.aborted) {
@ -431,6 +448,15 @@ export class OpenAiApi extends AiApi {
chunkCount++; chunkCount++;
finishReason = chunk.choices[0]?.finish_reason ?? finishReason; finishReason = chunk.choices[0]?.finish_reason ?? finishReason;
// Capture usage from the final streaming chunk
if (chunk.usage) {
usage = {
promptTokens: chunk.usage.prompt_tokens,
completionTokens: chunk.usage.completion_tokens,
};
}
const delta = chunk.choices[0]?.delta; const delta = chunk.choices[0]?.delta;
if (!delta) continue; if (!delta) continue;
@ -463,6 +489,7 @@ export class OpenAiApi extends AiApi {
chunkCount, chunkCount,
contentDeltaCount, contentDeltaCount,
toolDeltaCount, toolDeltaCount,
usage,
}; };
} }
@ -498,6 +525,10 @@ export class OpenAiApi extends AiApi {
const choice = response.choices[0]; const choice = response.choices[0];
const message = choice?.message; const message = choice?.message;
const content = typeof message?.content === "string" ? message.content : ""; const content = typeof message?.content === "string" ? message.content : "";
const usage = response.usage ? {
promptTokens: response.usage.prompt_tokens,
completionTokens: response.usage.completion_tokens,
} : undefined;
const assistantToolCalls = (message?.tool_calls || []) const assistantToolCalls = (message?.tool_calls || [])
.filter((tc) => tc.type === "function") .filter((tc) => tc.type === "function")
.map((tc) => ({ .map((tc) => ({
@ -524,6 +555,7 @@ export class OpenAiApi extends AiApi {
chunkCount: 0, chunkCount: 0,
contentDeltaCount: content ? 1 : 0, contentDeltaCount: content ? 1 : 0,
toolDeltaCount: assistantToolCalls.length, toolDeltaCount: assistantToolCalls.length,
usage,
}; };
} }
@ -592,7 +624,10 @@ export class OpenAiApi extends AiApi {
); );
} }
private buildStats(startTime: number): IAiChatResponse["stats"] { private buildStats(
startTime: number,
usage?: { promptTokens: number; completionTokens: number },
): IAiChatResponse["stats"] {
const seconds = (Date.now() - startTime) / 1000; const seconds = (Date.now() - startTime) / 1000;
return { return {
duration: { duration: {
@ -600,8 +635,8 @@ export class OpenAiApi extends AiApi {
text: numeral(seconds).format("hh:mm:ss"), text: numeral(seconds).format("hh:mm:ss"),
}, },
tokenCounts: { tokenCounts: {
input: 0, input: usage?.promptTokens ?? 0,
response: 0, response: usage?.completionTokens ?? 0,
thinking: 0, thinking: 0,
}, },
}; };

View File

@ -36,6 +36,7 @@ export interface IChatSession {
selectedModel: string; selectedModel: string;
reasoningEffort?: ReasoningEffort; reasoningEffort?: ReasoningEffort;
numCtx?: number; numCtx?: number;
contextWindowUsage: number;
stats: { stats: {
turnCount: number; turnCount: number;
toolCallCount: number; toolCallCount: number;

View File

@ -41,10 +41,19 @@ export type ToolCallMessage = (
response: string, response: string,
) => void; ) => void;
export interface IWorkOrderCompleteStats {
masterInputTokens: number;
masterOutputTokens: number;
masterThinkingTokens: number;
toolCallCount: number;
durationMs: number;
}
export type WorkOrderCompleteMessage = ( export type WorkOrderCompleteMessage = (
workOrderId: string, workOrderId: string,
success: boolean, success: boolean,
message?: string, message?: string,
stats?: IWorkOrderCompleteStats,
) => void; ) => void;
export type RequestCrashRecoveryMessage = (data: { export type RequestCrashRecoveryMessage = (data: {