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

Stream token usage on every agent message (#5342)

Browse Source
This commit is contained in:
David Katz
2025-10-29 23:23:25 -04:00
committed by GitHub
parent c875f13757
commit 37e1bb1d37
18 changed files with 214 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/desktop/openapi.json
+41
View File
@@ -3278,12 +3278,16 @@
"type": "object",
"required": [
"message",
"token_state",
"type"
],
"properties": {
"message": {
"$ref": "#/components/schemas/Message"
},
"token_state": {
"$ref": "#/components/schemas/TokenState"
},
"type": {
"type": "string",
"enum": [
@@ -4526,6 +4530,43 @@
}
}
},
"TokenState": {
"type": "object",
"required": [
"inputTokens",
"outputTokens",
"totalTokens",
"accumulatedInputTokens",
"accumulatedOutputTokens",
"accumulatedTotalTokens"
],
"properties": {
"accumulatedInputTokens": {
"type": "integer",
"format": "int32"
},
"accumulatedOutputTokens": {
"type": "integer",
"format": "int32"
},
"accumulatedTotalTokens": {
"type": "integer",
"format": "int32"
},
"inputTokens": {
"type": "integer",
"format": "int32"
},
"outputTokens": {
"type": "integer",
"format": "int32"
},
"totalTokens": {
"type": "integer",
"format": "int32"
}
}
},
"Tool": {
"type": "object",
"required": [
ui/desktop/src/api/types.gen.ts
+10
View File
@@ -367,6 +367,7 @@ export type MessageContent = (TextContent & {
export type MessageEvent = {
message: Message;
token_state: TokenState;
type: 'Message';
} | {
error: string;
@@ -789,6 +790,15 @@ export type ThinkingContent = {
thinking: string;
};
export type TokenState = {
accumulatedInputTokens: number;
accumulatedOutputTokens: number;
accumulatedTotalTokens: number;
inputTokens: number;
outputTokens: number;
totalTokens: number;
};
export type Tool = {
annotations?: ToolAnnotations | {
[key: string]: unknown;
ui/desktop/src/components/BaseChat.tsx
+8 -3
View File
@@ -132,6 +132,7 @@ function BaseChatContent({
sessionOutputTokens,
localInputTokens,
localOutputTokens,
tokenState,
commandHistory,
toolCallNotifications,
sessionMetadata,
@@ -442,9 +443,13 @@ function BaseChatContent({
commandHistory={commandHistory}
initialValue={input || ''}
setView={setView}
numTokens={sessionTokenCount}
inputTokens={sessionInputTokens || localInputTokens}
outputTokens={sessionOutputTokens || localOutputTokens}
totalTokens={tokenState?.totalTokens ?? sessionTokenCount}
accumulatedInputTokens={
tokenState?.accumulatedInputTokens ?? sessionInputTokens ?? localInputTokens
}
accumulatedOutputTokens={
tokenState?.accumulatedOutputTokens ?? sessionOutputTokens ?? localOutputTokens
}
droppedFiles={droppedFiles}
onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing
messages={messages}
ui/desktop/src/components/BaseChat2.tsx
+8 -3
View File
@@ -72,6 +72,7 @@ function BaseChatContent({
stopStreaming,
sessionLoadError,
setRecipeUserParams,
tokenState,
} = useChatStream({
sessionId,
onStreamFinish,
@@ -281,9 +282,13 @@ function BaseChatContent({
//commandHistory={commandHistory}
initialValue={initialPrompt}
setView={setView}
numTokens={session?.total_tokens || undefined}
inputTokens={session?.input_tokens || undefined}
outputTokens={session?.output_tokens || undefined}
totalTokens={tokenState?.totalTokens ?? session?.total_tokens ?? undefined}
accumulatedInputTokens={
tokenState?.accumulatedInputTokens ?? session?.accumulated_input_tokens ?? undefined
}
accumulatedOutputTokens={
tokenState?.accumulatedOutputTokens ?? session?.accumulated_output_tokens ?? undefined
}
droppedFiles={droppedFiles}
onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing
messages={messages}
ui/desktop/src/components/ChatInput.tsx
+12 -12
View File
@@ -70,9 +70,9 @@ interface ChatInputProps {
droppedFiles?: DroppedFile[];
onFilesProcessed?: () => void; // Callback to clear dropped files after processing
setView: (view: View) => void;
numTokens?: number;
inputTokens?: number;
outputTokens?: number;
totalTokens?: number;
accumulatedInputTokens?: number;
accumulatedOutputTokens?: number;
messages?: Message[];
sessionCosts?: {
[key: string]: {
@@ -103,9 +103,9 @@ export default function ChatInput({
droppedFiles = [],
onFilesProcessed,
setView,
numTokens,
inputTokens,
outputTokens,
totalTokens,
accumulatedInputTokens,
accumulatedOutputTokens,
messages = [],
disableAnimation = false,
sessionCosts,
@@ -505,16 +505,16 @@ export default function ChatInput({
clearAlerts();
// Show alert when either there is registered token usage, or we know the limit
if ((numTokens && numTokens > 0) || (isTokenLimitLoaded && tokenLimit)) {
if ((totalTokens && totalTokens > 0) || (isTokenLimitLoaded && tokenLimit)) {
addAlert({
type: AlertType.Info,
message: 'Context window',
progress: {
current: numTokens || 0,
current: totalTokens || 0,
total: tokenLimit,
},
showCompactButton: true,
compactButtonDisabled: !numTokens,
compactButtonDisabled: !totalTokens,
onCompact: () => {
window.dispatchEvent(new CustomEvent('hide-alert-popover'));
@@ -542,7 +542,7 @@ export default function ChatInput({
}
// We intentionally omit setView as it shouldn't trigger a re-render of alerts
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [numTokens, toolCount, tokenLimit, isTokenLimitLoaded, addAlert, clearAlerts]);
}, [totalTokens, toolCount, tokenLimit, isTokenLimitLoaded, addAlert, clearAlerts]);
// Cleanup effect for component unmount - prevent memory leaks
useEffect(() => {
@@ -1540,8 +1540,8 @@ export default function ChatInput({
<>
<div className="flex items-center h-full ml-1 mr-1">
<CostTracker
inputTokens={inputTokens}
outputTokens={outputTokens}
inputTokens={accumulatedInputTokens}
outputTokens={accumulatedOutputTokens}
sessionCosts={sessionCosts}
/>
</div>
ui/desktop/src/components/hub.tsx
+3 -3
View File
@@ -78,9 +78,9 @@ export default function Hub({
commandHistory={[]}
initialValue=""
setView={setView}
numTokens={0}
inputTokens={0}
outputTokens={0}
totalTokens={0}
accumulatedInputTokens={0}
accumulatedOutputTokens={0}
droppedFiles={[]}
onFilesProcessed={() => {}}
messages={[]}
ui/desktop/src/hooks/useChatEngine.ts
+2
View File
@@ -77,6 +77,7 @@ export const useChatEngine = ({
notifications,
session,
setError,
tokenState,
} = useMessageStream({
api: getApiUrl('/reply'),
id: chat.sessionId,
@@ -451,6 +452,7 @@ export const useChatEngine = ({
sessionOutputTokens,
localInputTokens,
localOutputTokens,
tokenState,
// UI helpers
commandHistory,
ui/desktop/src/hooks/useChatStream.ts
+15
View File
@@ -7,6 +7,7 @@ import {
reply,
resumeAgent,
Session,
TokenState,
updateFromSession,
updateSessionUserRecipeValues,
} from '../api';
@@ -60,6 +61,7 @@ interface UseChatStreamReturn {
setRecipeUserParams: (values: Record<string, string>) => Promise<void>;
stopStreaming: () => void;
sessionLoadError?: string;
tokenState: TokenState;
}
function pushMessage(currentMessages: Message[], incomingMsg: Message): Message[] {
@@ -88,6 +90,7 @@ async function streamFromResponse(
stream: AsyncIterable<MessageEvent>,
initialMessages: Message[],
updateMessages: (messages: Message[]) => void,
updateTokenState: (tokenState: TokenState) => void,
updateChatState: (state: ChatState) => void,
onFinish: (error?: string) => void
): Promise<void> {
@@ -119,6 +122,8 @@ async function streamFromResponse(
});
}
updateTokenState(event.token_state);
updateMessages(currentMessages);
break;
}
@@ -171,6 +176,14 @@ export function useChatStream({
const [session, setSession] = useState<Session>();
const [sessionLoadError, setSessionLoadError] = useState<string>();
const [chatState, setChatState] = useState<ChatState>(ChatState.Idle);
const [tokenState, setTokenState] = useState<TokenState>({
inputTokens: 0,
outputTokens: 0,
totalTokens: 0,
accumulatedInputTokens: 0,
accumulatedOutputTokens: 0,
accumulatedTotalTokens: 0,
});
const abortControllerRef = useRef<AbortController | null>(null);
useEffect(() => {
@@ -288,6 +301,7 @@ export function useChatStream({
stream,
currentMessages,
(messages: Message[]) => setMessagesAndLog(messages, 'streaming'),
setTokenState,
setChatState,
onFinish
);
@@ -373,5 +387,6 @@ export function useChatStream({
handleSubmit,
stopStreaming,
setRecipeUserParams,
tokenState,
};
}
ui/desktop/src/hooks/useMessageStream.ts
+16 -3
View File
@@ -6,7 +6,7 @@ import {
getCompactingMessage,
hasCompletedToolCalls,
} from '../types/message';
import { Conversation, Message, Role } from '../api';
import { Conversation, Message, Role, TokenState } from '../api';
import { getSession, Session } from '../api';
import { ChatState } from '../types/chatState';
@@ -35,7 +35,7 @@ export interface NotificationEvent {
// Event types for SSE stream
type MessageEvent =
| { type: 'Message'; message: Message }
| { type: 'Message'; message: Message; token_state: TokenState }
| { type: 'Error'; error: string }
| { type: 'Finish'; reason: string }
| { type: 'ModelChange'; model: string; mode: string }
@@ -165,6 +165,9 @@ export interface UseMessageStreamHelpers {
/** Clear error state */
setError: (error: Error | undefined) => void;
/** Real-time token state from server */
tokenState: TokenState;
}
/**
@@ -197,6 +200,14 @@ export function useMessageStream({
null
);
const [session, setSession] = useState<Session | null>(null);
const [tokenState, setTokenState] = useState<TokenState>({
inputTokens: 0,
outputTokens: 0,
totalTokens: 0,
accumulatedInputTokens: 0,
accumulatedOutputTokens: 0,
accumulatedTotalTokens: 0,
});
// expose a way to update the body so we can update the session id when CLE occurs
const updateMessageStreamBody = useCallback((newBody: object) => {
@@ -280,6 +291,8 @@ export function useMessageStream({
// Transition from waiting to streaming on first message
mutateChatState(ChatState.Streaming);
setTokenState(parsedEvent.token_state);
// Create a new message object with the properties preserved or defaulted
const newMessage: Message = {
...parsedEvent.message,
@@ -341,7 +354,6 @@ export function useMessageStream({
}
case 'UpdateConversation': {
currentMessages = parsedEvent.conversation;
setMessages(parsedEvent.conversation);
break;
}
@@ -650,5 +662,6 @@ export function useMessageStream({
currentModelInfo,
session,
setError,
tokenState,
};
}