tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

feat: Add comprehensive cost tracking display for LLM usage (#2992)

Browse Source 
Co-authored-by: jack <jack@deck.local>
Co-authored-by: Bradley Axen <baxen@squareup.com>
This commit is contained in:
jack
2025-06-26 10:46:14 +01:00
committed by GitHub
parent 9d48ed980f
commit d2ff4f3746
24 changed files with 1985 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/desktop/src/hooks/useMessageStream.ts
+52 -3
View File
@@ -2,12 +2,26 @@ import { useState, useCallback, useEffect, useRef, useId } from 'react';
import useSWR from 'swr';
import { getSecretKey } from '../config';
import { Message, createUserMessage, hasCompletedToolCalls } from '../types/message';
import { getSessionHistory } from '../api';
// Ensure TextDecoder is available in the global scope
const TextDecoder = globalThis.TextDecoder;
type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue };
export interface SessionMetadata {
workingDir: string;
description: string;
scheduleId: string | null;
messageCount: number;
totalTokens: number | null;
inputTokens: number | null;
outputTokens: number | null;
accumulatedTotalTokens: number | null;
accumulatedInputTokens: number | null;
accumulatedOutputTokens: number | null;
}
export interface NotificationEvent {
type: 'Notification';
request_id: string;
@@ -141,9 +155,12 @@ export interface UseMessageStreamHelpers {
updateMessageStreamBody?: (newBody: object) => void;
notifications: NotificationEvent[];
/** Current model info from the backend */
currentModelInfo: { model: string; mode: string } | null;
/** Session metadata including token counts */
sessionMetadata: SessionMetadata | null;
}
/**
@@ -172,7 +189,10 @@ export function useMessageStream({
});
const [notifications, setNotifications] = useState<NotificationEvent[]>([]);
const [currentModelInfo, setCurrentModelInfo] = useState<{ model: string; mode: string } | null>(null);
const [currentModelInfo, setCurrentModelInfo] = useState<{ model: string; mode: string } | null>(
null
);
const [sessionMetadata, setSessionMetadata] = useState<SessionMetadata | null>(null);
// expose a way to update the body so we can update the session id when CLE occurs
const updateMessageStreamBody = useCallback((newBody: object) => {
@@ -291,13 +311,41 @@ export function useMessageStream({
case 'Error':
throw new Error(parsedEvent.error);
case 'Finish':
case 'Finish': {
// Call onFinish with the last message if available
if (onFinish && currentMessages.length > 0) {
const lastMessage = currentMessages[currentMessages.length - 1];
onFinish(lastMessage, parsedEvent.reason);
}
// Fetch updated session metadata with token counts
const sessionId = (extraMetadataRef.current.body as Record<string, unknown>)?.session_id as string;
if (sessionId) {
try {
const sessionResponse = await getSessionHistory({
path: { session_id: sessionId },
});
if (sessionResponse.data?.metadata) {
setSessionMetadata({
workingDir: sessionResponse.data.metadata.working_dir,
description: sessionResponse.data.metadata.description,
scheduleId: sessionResponse.data.metadata.schedule_id || null,
messageCount: sessionResponse.data.metadata.message_count,
totalTokens: sessionResponse.data.metadata.total_tokens || null,
inputTokens: sessionResponse.data.metadata.input_tokens || null,
outputTokens: sessionResponse.data.metadata.output_tokens || null,
accumulatedTotalTokens: sessionResponse.data.metadata.accumulated_total_tokens || null,
accumulatedInputTokens: sessionResponse.data.metadata.accumulated_input_tokens || null,
accumulatedOutputTokens: sessionResponse.data.metadata.accumulated_output_tokens || null,
});
}
} catch (error) {
console.error('Failed to fetch session metadata:', error);
}
}
break;
}
}
} catch (e) {
console.error('Error parsing SSE event:', e);
@@ -559,5 +607,6 @@ export function useMessageStream({
updateMessageStreamBody,
notifications,
currentModelInfo,
sessionMetadata,
};
}