zen: fix parsing cache write tokens

2 files changed, 8 insertions, 6 deletions

diff --git a/packages/console/app/src/routes/zen/handler.ts b/packages/console/app/src/routes/zen/handler.ts
index 2ff3ca073..deab024c7 100644
--- a/packages/console/app/src/routes/zen/handler.ts
+++ b/packages/console/app/src/routes/zen/handler.ts
@@ -161,7 +161,6 @@ export async function handler(
                 logger.metric({ time_to_first_byte: Date.now() - startTimestamp })
               }
               responseLength += value.length
-              console.log(decoder.decode(value, { stream: true }))
               buffer += decoder.decode(value, { stream: true })
 
               const parts = buffer.split("\n\n")
diff --git a/packages/console/app/src/routes/zen/v1/messages.ts b/packages/console/app/src/routes/zen/v1/messages.ts
index 1fd85d5c7..4a7dda5f7 100644
--- a/packages/console/app/src/routes/zen/v1/messages.ts
+++ b/packages/console/app/src/routes/zen/v1/messages.ts
@@ -30,22 +30,25 @@ export function POST(input: APIEvent) {
 
       let json
       try {
-        json = JSON.parse(data.slice(6)) as { usage?: Usage }
+        json = JSON.parse(data.slice(6))
       } catch (e) {
         return
       }
 
-      if (!json.usage) return
+      // ie. { type: "message_start"; message: { usage: Usage } }
+      // ie. { type: "message_delta"; usage: Usage }
+      const usageUpdate = json.usage ?? json.message?.usage
+      if (!usageUpdate) return
       usage = {
         ...usage,
-        ...json.usage,
+        ...usageUpdate,
         cache_creation: {
           ...usage?.cache_creation,
-          ...json.usage.cache_creation,
+          ...usageUpdate.cache_creation,
         },
         server_tool_use: {
           ...usage?.server_tool_use,
-          ...json.usage.server_tool_use,
+          ...usageUpdate.server_tool_use,
         },
       }
     },