dtp/gadget
dtp/gadget
1
0
Fork 0
Code Actions Packages Projects Releases Activity

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

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
gadget-code/frontend/src/components/SessionPanel.tsx
View File

@ -41,6 +41,7 @@ interface SessionPanelProps {
truncateModelName: (name: string, maxLength?: number) => string;
getSelectedModelCapabilities: () => { hasThinking: boolean } | null;
getSelectedModel: () => AiModel | null;
contextWindowUsage: number;
}
export default function SessionPanel({
@ -82,7 +83,9 @@ export default function SessionPanel({
truncateModelName,
getSelectedModelCapabilities,
getSelectedModel,
contextWindowUsage,
}: SessionPanelProps) {
const usagePercent = sessionNumCtx > 0 ? (contextWindowUsage / sessionNumCtx) * 100 : 0;
return (
<div className="border-b border-border-subtle shrink-0">
<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 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>
);

9
gadget-code/frontend/src/lib/api.ts
View File

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

37
gadget-code/frontend/src/pages/ChatSessionView.tsx
View File

@ -2,7 +2,7 @@ import { useState, useEffect, useRef, useContext, useCallback } from 'react';
import { useParams, useNavigate, useLocation } from 'react-router-dom';
import { BrushCleaning, Maximize2 } from 'lucide-react';
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 SessionPanel from '../components/SessionPanel';
import ProjectPanel from '../components/ProjectPanel';
@ -686,7 +686,7 @@ export default function ChatSessionView() {
}, 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
// Just clean up frontend streaming state and update status
if (streamingStateRef.current.has(turnId)) {
@ -734,6 +734,38 @@ export default function ChatSessionView() {
setIsAborting(false);
setIsProcessing(false);
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]);
const handleWorkspaceModeChanged = useCallback((mode: string) => {
@ -1441,6 +1473,7 @@ export default function ChatSessionView() {
truncateModelName={truncateModelName}
getSelectedModelCapabilities={getSelectedModelCapabilities}
getSelectedModel={getSelectedModel}
contextWindowUsage={session?.contextWindowUsage ?? 0}
/>
<ProjectPanel

12
gadget-code/src/lib/code-session.ts
View File

@ -23,6 +23,7 @@ import {
SubmitPromptCallback,
AbortWorkOrderCallback,
SubProcessStat,
IWorkOrderCompleteStats,
} from "@gadget/api";
import ChatSession from "../models/chat-session.ts";
@ -669,7 +670,16 @@ export class CodeSession extends SocketSession {
turnId: string,
success: boolean,
message?: string,
stats?: IWorkOrderCompleteStats,
): 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;
}
}
}

56
gadget-code/src/lib/drone-session.ts
View File

@ -12,6 +12,7 @@ import {
WorkspaceMode,
IChatToolCall,
IChatTurnBlock,
IWorkOrderCompleteStats,
} from "@gadget/api";
import {
GadgetSocket,
@ -20,6 +21,7 @@ import {
} from "./socket-session.js";
import { SocketService } from "../services/index.js";
import { ChatTurn } from "../models/chat-turn.js";
import ChatSession from "../models/chat-session.js";
import MessageQueue, { type QueuedMessage } from "./message-queue.js";
interface IStreamingBuffer {
@ -292,6 +294,7 @@ export class DroneSession extends SocketSession {
turnId: string,
success: boolean,
message?: string,
stats?: IWorkOrderCompleteStats,
): Promise<void> {
if (!this.chatSessionId) {
this.log.warn("workOrderComplete event received but no chat session is active");
@ -314,20 +317,55 @@ export class DroneSession extends SocketSession {
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();
// 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(
this.chatSessionId,
);
codeSession.onWorkOrderComplete(turnId, success, message);
codeSession.onWorkOrderComplete(turnId, success, message, stats);
this.currentTurnId = undefined;
} catch (error) {
// Routing failed - queue to Redis
await MessageQueue.enqueue(this.chatSessionId, {
type: 'workOrderComplete',
args: [turnId, success, message],
args: [turnId, success, message, stats],
timestamp: Date.now(),
});
this.log.debug("queued workOrderComplete message", { chatSessionId: this.chatSessionId });
@ -343,6 +381,17 @@ export class DroneSession extends SocketSession {
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.
*/
@ -669,7 +718,8 @@ export class DroneSession extends SocketSession {
codeSession.onWorkOrderComplete(
msg.args[0] as string,
msg.args[1] as boolean,
msg.args[2] as string,
msg.args[2] as string | undefined,
msg.args[3] as IWorkOrderCompleteStats | undefined,
);
break;
case 'log':

1
gadget-code/src/models/chat-session.ts
View File

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

1
gadget-code/src/services/chat-session.ts
View File

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

3
gadget-drone/src/services/agent.test.ts
View File

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

55
gadget-drone/src/services/agent.ts
View File

@ -58,6 +58,14 @@ import {
ReadSkillTool,
type DroneToolboxEnvironment,
} from "../tools/index.ts";
interface IWorkOrderCompleteStats {
masterInputTokens: number;
masterOutputTokens: number;
masterThinkingTokens: number;
toolCallCount: number;
durationMs: number;
}
import type { IChatExportData } from "@gadget/ai-toolbox";
export interface IAgentWorkOrder {
@ -164,8 +172,10 @@ class AgentService extends GadgetService {
const { turn } = workOrder;
let toolCallCount = 0;
let inputTokens = 0;
let outputTokens = 0;
let masterInputTokens = 0;
let masterOutputTokens = 0;
let masterThinkingTokens = 0;
const startTime = Date.now();
// Build the full message array that grows with each iteration
const messages: IContextChatMessage[] = [];
@ -258,6 +268,11 @@ class AgentService extends GadgetService {
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
if (response.toolCalls && response.toolCalls.length > 0) {
continueLoop = true;
@ -338,17 +353,30 @@ class AgentService extends GadgetService {
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) {
if (cause instanceof Error && cause.name === "AbortError") {
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;
}
const msg = cause instanceof Error ? cause.message : String(cause);
@ -356,7 +384,14 @@ class AgentService extends GadgetService {
turnId: turn._id,
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;
} finally {
this.abortController = null;
@ -740,11 +775,11 @@ class AgentService extends GadgetService {
inputTokens += response.stats?.tokenCounts?.input ?? 0;
if (response.thinking) {
thinkingTokenCount += Math.ceil(response.thinking.length / 4);
thinkingTokenCount += response.stats?.tokenCounts?.thinking ?? 0;
}
if (response.response) {
fullResponse += response.response;
outputTokens += Math.ceil(response.response.length / 4);
outputTokens += response.stats?.tokenCounts?.response ?? 0;
}
if (response.thinking) {
fullThinking = (fullThinking + response.thinking).trim();
@ -800,8 +835,6 @@ class AgentService extends GadgetService {
toolName: toolCall.function.name,
});
inputTokens += Math.ceil(toolArgsRaw.length / 4);
outputTokens += Math.ceil(result.length / 4);
}
}
}

49
packages/ai/src/openai.ts
View File

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

1
packages/api/src/interfaces/chat-session.ts
View File

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

9
packages/api/src/messages/drone.ts
View File

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