From f19a9aac6674bd4c5e582fd14ca41755061d78f1 Mon Sep 17 00:00:00 2001 From: David Katz Date: Thu, 6 Nov 2025 11:19:36 -0500 Subject: [PATCH 1/6] show token state on session load --- ui/desktop/src/components/BaseChat.tsx | 6 +++--- ui/desktop/src/hooks/useAgent.ts | 16 ++++++++++++++++ ui/desktop/src/hooks/useChatEngine.ts | 13 ++++++++++++- ui/desktop/src/types/chat.ts | 3 ++- 4 files changed, 33 insertions(+), 5 deletions(-) diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index 67090686db2c..fa96740f064c 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -443,12 +443,12 @@ function BaseChatContent({ commandHistory={commandHistory} initialValue={input || ''} setView={setView} - totalTokens={tokenState?.totalTokens ?? sessionTokenCount} + totalTokens={tokenState?.totalTokens || sessionTokenCount} accumulatedInputTokens={ - tokenState?.accumulatedInputTokens ?? sessionInputTokens ?? localInputTokens + tokenState?.accumulatedInputTokens || sessionInputTokens || localInputTokens } accumulatedOutputTokens={ - tokenState?.accumulatedOutputTokens ?? sessionOutputTokens ?? localOutputTokens + tokenState?.accumulatedOutputTokens || sessionOutputTokens || localOutputTokens } droppedFiles={droppedFiles} onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing diff --git a/ui/desktop/src/hooks/useAgent.ts b/ui/desktop/src/hooks/useAgent.ts index 136554ab4cee..027ba2a6a0ef 100644 --- a/ui/desktop/src/hooks/useAgent.ts +++ b/ui/desktop/src/hooks/useAgent.ts @@ -87,6 +87,14 @@ export function useAgent(): UseAgentReturn { messages, recipe: agentSession.recipe, recipeParameterValues: agentSession.user_recipe_values || null, + tokenState: { + totalTokens: agentSession.total_tokens ?? 0, + inputTokens: agentSession.input_tokens ?? 0, + outputTokens: agentSession.output_tokens ?? 0, + accumulatedTotalTokens: agentSession.accumulated_total_tokens ?? 0, + accumulatedInputTokens: agentSession.accumulated_input_tokens ?? 0, + accumulatedOutputTokens: agentSession.accumulated_output_tokens ?? 0, + }, }; } @@ -189,6 +197,14 @@ export function useAgent(): UseAgentReturn { messages: messages, recipe: recipe, recipeParameterValues: agentSession.user_recipe_values || null, + tokenState: { + totalTokens: agentSession.total_tokens ?? 0, + inputTokens: agentSession.input_tokens ?? 0, + outputTokens: agentSession.output_tokens ?? 0, + accumulatedTotalTokens: agentSession.accumulated_total_tokens ?? 0, + accumulatedInputTokens: agentSession.accumulated_input_tokens ?? 0, + accumulatedOutputTokens: agentSession.accumulated_output_tokens ?? 0, + }, }; setAgentState(AgentState.INITIALIZED); diff --git a/ui/desktop/src/hooks/useChatEngine.ts b/ui/desktop/src/hooks/useChatEngine.ts index 7c60038e6316..4d8204057f43 100644 --- a/ui/desktop/src/hooks/useChatEngine.ts +++ b/ui/desktop/src/hooks/useChatEngine.ts @@ -36,6 +36,18 @@ export const useChatEngine = ({ const [localOutputTokens, setLocalOutputTokens] = useState(0); const [powerSaveTimeoutId, setPowerSaveTimeoutId] = useState(null); + // Initialize token state from chat when session is first loaded or changes + useEffect(() => { + if ( + chat.tokenState?.accumulatedTotalTokens !== undefined && + chat.tokenState?.accumulatedTotalTokens !== null + ) { + setSessionTokenCount(chat.tokenState.totalTokens || 0); + setSessionInputTokens(chat.tokenState.accumulatedInputTokens || 0); + setSessionOutputTokens(chat.tokenState.accumulatedOutputTokens || 0); + } + }, [chat.sessionId, chat.tokenState]); + // Track pending edited message const [pendingEdit, setPendingEdit] = useState<{ id: string; content: string } | null>(null); @@ -213,7 +225,6 @@ export const useChatEngine = ({ // Update token counts when session changes from the message stream useEffect(() => { - console.log('Session received:', session); if (session) { setSessionTokenCount(session.total_tokens || 0); setSessionInputTokens(session.accumulated_input_tokens || 0); diff --git a/ui/desktop/src/types/chat.ts b/ui/desktop/src/types/chat.ts index 77d875c1fb6b..636a366ac7e4 100644 --- a/ui/desktop/src/types/chat.ts +++ b/ui/desktop/src/types/chat.ts @@ -1,5 +1,5 @@ import { Recipe } from '../recipe'; -import { Message } from '../api'; +import { Message, TokenState } from '../api'; export interface ChatType { sessionId: string; @@ -9,4 +9,5 @@ export interface ChatType { recipe?: Recipe | null; // Add recipe configuration to chat state resolvedRecipe?: Recipe | null; // Add resolved recipe with parameter values rendered to chat state recipeParameterValues?: Record | null; // Add recipe parameters to chat state + tokenState?: TokenState | null; // Token state from session } From fd801aec3cd01bda471d8ee0a0c1540130896148 Mon Sep 17 00:00:00 2001 From: David Katz Date: Thu, 6 Nov 2025 11:26:23 -0500 Subject: [PATCH 2/6] simplify fix --- ui/desktop/src/hooks/useAgent.ts | 16 ---------------- ui/desktop/src/hooks/useChatEngine.ts | 12 ------------ ui/desktop/src/types/chat.ts | 3 +-- 3 files changed, 1 insertion(+), 30 deletions(-) diff --git a/ui/desktop/src/hooks/useAgent.ts b/ui/desktop/src/hooks/useAgent.ts index 027ba2a6a0ef..136554ab4cee 100644 --- a/ui/desktop/src/hooks/useAgent.ts +++ b/ui/desktop/src/hooks/useAgent.ts @@ -87,14 +87,6 @@ export function useAgent(): UseAgentReturn { messages, recipe: agentSession.recipe, recipeParameterValues: agentSession.user_recipe_values || null, - tokenState: { - totalTokens: agentSession.total_tokens ?? 0, - inputTokens: agentSession.input_tokens ?? 0, - outputTokens: agentSession.output_tokens ?? 0, - accumulatedTotalTokens: agentSession.accumulated_total_tokens ?? 0, - accumulatedInputTokens: agentSession.accumulated_input_tokens ?? 0, - accumulatedOutputTokens: agentSession.accumulated_output_tokens ?? 0, - }, }; } @@ -197,14 +189,6 @@ export function useAgent(): UseAgentReturn { messages: messages, recipe: recipe, recipeParameterValues: agentSession.user_recipe_values || null, - tokenState: { - totalTokens: agentSession.total_tokens ?? 0, - inputTokens: agentSession.input_tokens ?? 0, - outputTokens: agentSession.output_tokens ?? 0, - accumulatedTotalTokens: agentSession.accumulated_total_tokens ?? 0, - accumulatedInputTokens: agentSession.accumulated_input_tokens ?? 0, - accumulatedOutputTokens: agentSession.accumulated_output_tokens ?? 0, - }, }; setAgentState(AgentState.INITIALIZED); diff --git a/ui/desktop/src/hooks/useChatEngine.ts b/ui/desktop/src/hooks/useChatEngine.ts index 4d8204057f43..19a333aefda9 100644 --- a/ui/desktop/src/hooks/useChatEngine.ts +++ b/ui/desktop/src/hooks/useChatEngine.ts @@ -36,18 +36,6 @@ export const useChatEngine = ({ const [localOutputTokens, setLocalOutputTokens] = useState(0); const [powerSaveTimeoutId, setPowerSaveTimeoutId] = useState(null); - // Initialize token state from chat when session is first loaded or changes - useEffect(() => { - if ( - chat.tokenState?.accumulatedTotalTokens !== undefined && - chat.tokenState?.accumulatedTotalTokens !== null - ) { - setSessionTokenCount(chat.tokenState.totalTokens || 0); - setSessionInputTokens(chat.tokenState.accumulatedInputTokens || 0); - setSessionOutputTokens(chat.tokenState.accumulatedOutputTokens || 0); - } - }, [chat.sessionId, chat.tokenState]); - // Track pending edited message const [pendingEdit, setPendingEdit] = useState<{ id: string; content: string } | null>(null); diff --git a/ui/desktop/src/types/chat.ts b/ui/desktop/src/types/chat.ts index 636a366ac7e4..77d875c1fb6b 100644 --- a/ui/desktop/src/types/chat.ts +++ b/ui/desktop/src/types/chat.ts @@ -1,5 +1,5 @@ import { Recipe } from '../recipe'; -import { Message, TokenState } from '../api'; +import { Message } from '../api'; export interface ChatType { sessionId: string; @@ -9,5 +9,4 @@ export interface ChatType { recipe?: Recipe | null; // Add recipe configuration to chat state resolvedRecipe?: Recipe | null; // Add resolved recipe with parameter values rendered to chat state recipeParameterValues?: Record | null; // Add recipe parameters to chat state - tokenState?: TokenState | null; // Token state from session } From a36c3fd8eb316810cb9581be029d22915776c541 Mon Sep 17 00:00:00 2001 From: David Katz Date: Thu, 6 Nov 2025 11:44:17 -0500 Subject: [PATCH 3/6] fix token state --- crates/goose-server/src/routes/reply.rs | 52 +++++++++++++----------- ui/desktop/src/hooks/useMessageStream.ts | 4 +- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 6c5425d9a04e..f5822c23c430 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -133,6 +133,7 @@ pub enum MessageEvent { }, Finish { reason: String, + token_state: TokenState, }, ModelChange { model: String, @@ -149,6 +150,30 @@ pub enum MessageEvent { Ping, } +async fn get_token_state(session_id: &str) -> TokenState { + match SessionManager::get_session(session_id, false).await { + Ok(session) => TokenState { + input_tokens: session.input_tokens.unwrap_or(0), + output_tokens: session.output_tokens.unwrap_or(0), + total_tokens: session.total_tokens.unwrap_or(0), + accumulated_input_tokens: session.accumulated_input_tokens.unwrap_or(0), + accumulated_output_tokens: session.accumulated_output_tokens.unwrap_or(0), + accumulated_total_tokens: session.accumulated_total_tokens.unwrap_or(0), + }, + Err(e) => { + tracing::warn!("Failed to fetch session token state for {}: {}", session_id, e); + TokenState { + input_tokens: 0, + output_tokens: 0, + total_tokens: 0, + accumulated_input_tokens: 0, + accumulated_output_tokens: 0, + accumulated_total_tokens: 0, + } + } + } +} + async fn stream_event( event: MessageEvent, tx: &mpsc::Sender, @@ -321,29 +346,7 @@ pub async fn reply( all_messages.push(message.clone()); - let token_state = match SessionManager::get_session(&session_id, false).await { - Ok(session) => { - TokenState { - input_tokens: session.input_tokens.unwrap_or(0), - output_tokens: session.output_tokens.unwrap_or(0), - total_tokens: session.total_tokens.unwrap_or(0), - accumulated_input_tokens: session.accumulated_input_tokens.unwrap_or(0), - accumulated_output_tokens: session.accumulated_output_tokens.unwrap_or(0), - accumulated_total_tokens: session.accumulated_total_tokens.unwrap_or(0), - } - }, - Err(e) => { - tracing::warn!("Failed to fetch session for token state: {}", e); - TokenState { - input_tokens: 0, - output_tokens: 0, - total_tokens: 0, - accumulated_input_tokens: 0, - accumulated_output_tokens: 0, - accumulated_total_tokens: 0, - } - } - }; + let token_state = get_token_state(&session_id).await; stream_event(MessageEvent::Message { message, token_state }, &tx, &cancel_token).await; } @@ -437,9 +440,12 @@ pub async fn reply( ); } + let final_token_state = get_token_state(&session_id).await; + let _ = stream_event( MessageEvent::Finish { reason: "stop".to_string(), + token_state: final_token_state, }, &task_tx, &cancel_token, diff --git a/ui/desktop/src/hooks/useMessageStream.ts b/ui/desktop/src/hooks/useMessageStream.ts index 2eb98835935b..4e36c7d9ec76 100644 --- a/ui/desktop/src/hooks/useMessageStream.ts +++ b/ui/desktop/src/hooks/useMessageStream.ts @@ -37,7 +37,7 @@ export interface NotificationEvent { type MessageEvent = | { type: 'Message'; message: Message; token_state: TokenState } | { type: 'Error'; error: string } - | { type: 'Finish'; reason: string } + | { type: 'Finish'; reason: string; token_state: TokenState } | { type: 'ModelChange'; model: string; mode: string } | { type: 'UpdateConversation'; conversation: Conversation } | NotificationEvent; @@ -368,6 +368,8 @@ export function useMessageStream({ } case 'Finish': { + setTokenState(parsedEvent.token_state); + if (onFinish && currentMessages.length > 0) { const lastMessage = currentMessages[currentMessages.length - 1]; onFinish(lastMessage, parsedEvent.reason); From a28c5db5219fae6f5019f222a22415fdc228b071 Mon Sep 17 00:00:00 2001 From: David Katz Date: Thu, 6 Nov 2025 11:45:35 -0500 Subject: [PATCH 4/6] fmt --- crates/goose-server/src/routes/reply.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index f5822c23c430..0e9a7baad746 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -161,7 +161,11 @@ async fn get_token_state(session_id: &str) -> TokenState { accumulated_total_tokens: session.accumulated_total_tokens.unwrap_or(0), }, Err(e) => { - tracing::warn!("Failed to fetch session token state for {}: {}", session_id, e); + tracing::warn!( + "Failed to fetch session token state for {}: {}", + session_id, + e + ); TokenState { input_tokens: 0, output_tokens: 0, From c4009486a34716cf325ce0f9b954561f0f7bbbaf Mon Sep 17 00:00:00 2001 From: David Katz Date: Thu, 6 Nov 2025 12:11:42 -0500 Subject: [PATCH 5/6] Default --- crates/goose-server/src/routes/reply.rs | 21 +++++++-------------- crates/goose/src/conversation/message.rs | 2 +- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 0e9a7baad746..5a340fa3d947 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -151,31 +151,24 @@ pub enum MessageEvent { } async fn get_token_state(session_id: &str) -> TokenState { - match SessionManager::get_session(session_id, false).await { - Ok(session) => TokenState { + SessionManager::get_session(session_id, false) + .await + .map(|session| TokenState { input_tokens: session.input_tokens.unwrap_or(0), output_tokens: session.output_tokens.unwrap_or(0), total_tokens: session.total_tokens.unwrap_or(0), accumulated_input_tokens: session.accumulated_input_tokens.unwrap_or(0), accumulated_output_tokens: session.accumulated_output_tokens.unwrap_or(0), accumulated_total_tokens: session.accumulated_total_tokens.unwrap_or(0), - }, - Err(e) => { + }) + .inspect_err(|e| { tracing::warn!( "Failed to fetch session token state for {}: {}", session_id, e ); - TokenState { - input_tokens: 0, - output_tokens: 0, - total_tokens: 0, - accumulated_input_tokens: 0, - accumulated_output_tokens: 0, - accumulated_total_tokens: 0, - } - } - } + }) + .unwrap_or_default() } async fn stream_event( diff --git a/crates/goose/src/conversation/message.rs b/crates/goose/src/conversation/message.rs index cc7d161dd841..2f18d038836f 100644 --- a/crates/goose/src/conversation/message.rs +++ b/crates/goose/src/conversation/message.rs @@ -711,7 +711,7 @@ impl Message { } } -#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)] #[serde(rename_all = "camelCase")] pub struct TokenState { pub input_tokens: i32, From 24875ed5f799a498287e2381d124ff9eae03ebce Mon Sep 17 00:00:00 2001 From: David Katz Date: Thu, 6 Nov 2025 12:44:31 -0500 Subject: [PATCH 6/6] gen openapi --- ui/desktop/openapi.json | 4 ++++ ui/desktop/src/api/types.gen.ts | 1 + 2 files changed, 5 insertions(+) diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index d7412d5bf626..7622534994d9 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -3321,12 +3321,16 @@ "type": "object", "required": [ "reason", + "token_state", "type" ], "properties": { "reason": { "type": "string" }, + "token_state": { + "$ref": "#/components/schemas/TokenState" + }, "type": { "type": "string", "enum": [ diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 6468fd3ee3cb..098489445a7b 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -374,6 +374,7 @@ export type MessageEvent = { type: 'Error'; } | { reason: string; + token_state: TokenState; type: 'Finish'; } | { mode: string;