From 4efe8a539dc6dd4eaffd0e4d2d97ef205050151e Mon Sep 17 00:00:00 2001
From: David Katz <dkatz@squareup.com>
Date: Fri, 7 Nov 2025 15:14:30 -0500
Subject: [PATCH] new pipethru mechanism

---
 crates/goose-server/src/routes/reply.rs  |  2 ++
 crates/goose/src/conversation/message.rs |  2 ++
 ui/desktop/openapi.json                  | 12 +++++++++++-
 ui/desktop/src/api/types.gen.ts          |  2 ++
 ui/desktop/src/components/BaseChat.tsx   |  6 ++++--
 ui/desktop/src/hooks/useChatStream.ts    |  2 ++
 ui/desktop/src/hooks/useMessageStream.ts |  2 ++
 7 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs
index 5a340fa3d9..011c285bb7 100644
--- a/crates/goose-server/src/routes/reply.rs
+++ b/crates/goose-server/src/routes/reply.rs
@@ -160,6 +160,8 @@ async fn get_token_state(session_id: &str) -> TokenState {
             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),
+            accumulated_cache_read_input_tokens: session.accumulated_cache_read_input_tokens.unwrap_or(0),
+            accumulated_cache_write_input_tokens: session.accumulated_cache_write_input_tokens.unwrap_or(0),
         })
         .inspect_err(|e| {
             tracing::warn!(
diff --git a/crates/goose/src/conversation/message.rs b/crates/goose/src/conversation/message.rs
index 2f18d03883..ef918163ff 100644
--- a/crates/goose/src/conversation/message.rs
+++ b/crates/goose/src/conversation/message.rs
@@ -720,6 +720,8 @@ pub struct TokenState {
     pub accumulated_input_tokens: i32,
     pub accumulated_output_tokens: i32,
     pub accumulated_total_tokens: i32,
+    pub accumulated_cache_read_input_tokens: i32,
+    pub accumulated_cache_write_input_tokens: i32,
 }
 
 #[cfg(test)]
diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json
index 5182b0be31..d383bac7bc 100644
--- a/ui/desktop/openapi.json
+++ b/ui/desktop/openapi.json
@@ -4556,9 +4556,19 @@
           "totalTokens",
           "accumulatedInputTokens",
           "accumulatedOutputTokens",
-          "accumulatedTotalTokens"
+          "accumulatedTotalTokens",
+          "accumulatedCacheReadInputTokens",
+          "accumulatedCacheWriteInputTokens"
         ],
         "properties": {
+          "accumulatedCacheReadInputTokens": {
+            "type": "integer",
+            "format": "int32"
+          },
+          "accumulatedCacheWriteInputTokens": {
+            "type": "integer",
+            "format": "int32"
+          },
           "accumulatedInputTokens": {
             "type": "integer",
             "format": "int32"
diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts
index 2f00953f7e..ef858273b8 100644
--- a/ui/desktop/src/api/types.gen.ts
+++ b/ui/desktop/src/api/types.gen.ts
@@ -795,6 +795,8 @@ export type ThinkingContent = {
 };
 
 export type TokenState = {
+    accumulatedCacheReadInputTokens: number;
+    accumulatedCacheWriteInputTokens: number;
     accumulatedInputTokens: number;
     accumulatedOutputTokens: number;
     accumulatedTotalTokens: number;
diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx
index fe7d683bab..4276b24f30 100644
--- a/ui/desktop/src/components/BaseChat.tsx
+++ b/ui/desktop/src/components/BaseChat.tsx
@@ -456,8 +456,10 @@ function BaseChatContent({
             accumulatedOutputTokens={
               tokenState?.accumulatedOutputTokens || sessionOutputTokens || localOutputTokens
             }
-            cacheReadTokens={sessionCacheReadTokens}
-            cacheWriteTokens={sessionCacheWriteTokens}
+            cacheReadTokens={tokenState?.accumulatedCacheReadInputTokens || sessionCacheReadTokens}
+            cacheWriteTokens={
+              tokenState?.accumulatedCacheWriteInputTokens || sessionCacheWriteTokens
+            }
             droppedFiles={droppedFiles}
             onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing
             messages={messages}
diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts
index e084b04d6e..8fba64e1d9 100644
--- a/ui/desktop/src/hooks/useChatStream.ts
+++ b/ui/desktop/src/hooks/useChatStream.ts
@@ -185,6 +185,8 @@ export function useChatStream({
     accumulatedInputTokens: 0,
     accumulatedOutputTokens: 0,
     accumulatedTotalTokens: 0,
+    accumulatedCacheReadInputTokens: 0,
+    accumulatedCacheWriteInputTokens: 0,
   });
   const abortControllerRef = useRef<AbortController | null>(null);
 
diff --git a/ui/desktop/src/hooks/useMessageStream.ts b/ui/desktop/src/hooks/useMessageStream.ts
index 4e36c7d9ec..57e4cee037 100644
--- a/ui/desktop/src/hooks/useMessageStream.ts
+++ b/ui/desktop/src/hooks/useMessageStream.ts
@@ -207,6 +207,8 @@ export function useMessageStream({
     accumulatedInputTokens: 0,
     accumulatedOutputTokens: 0,
     accumulatedTotalTokens: 0,
+    accumulatedCacheReadInputTokens: 0,
+    accumulatedCacheWriteInputTokens: 0,
   });
 
   // expose a way to update the body so we can update the session id when CLE occurs