wip: zen

3 files changed, 82 insertions, 10 deletions

diff --git a/cloud/app/src/routes/zen/v1/chat/completions.ts b/cloud/app/src/routes/zen/v1/chat/completions.ts
index 0a3b0e759..b68264cc4 100644
--- a/cloud/app/src/routes/zen/v1/chat/completions.ts
+++ b/cloud/app/src/routes/zen/v1/chat/completions.ts
@@ -2,14 +2,19 @@ import type { APIEvent } from "@solidjs/start/server"
 import { handler } from "~/util/zen"
 
 export function POST(input: APIEvent) {
+  let usage: any
   return handler(input, {
-    transformBody: (body: any) => ({
+    modifyBody: (body: any) => ({
       ...body,
       stream_options: {
         include_usage: true,
       },
     }),
-    parseUsageChunk: (chunk: string) => {
+    setAuthHeader: (headers: Headers, apiKey: string) => {
+      headers.set("authorization", `Bearer ${apiKey}`)
+    },
+    parseApiKey: (headers: Headers) => headers.get("authorization")?.split(" ")[1],
+    onStreamPart: (chunk: string) => {
       if (!chunk.startsWith("data: ")) return
 
       let json
@@ -19,15 +24,15 @@ export function POST(input: APIEvent) {
         return
       }
 
-      return json.usage
+      if (!json.usage) return
+      usage = json.usage
     },
-    buildUsage: (usage: any) => ({
+    getStreamUsage: () => usage,
+    normalizeUsage: (usage: any) => ({
       inputTokens: usage.prompt_tokens ?? 0,
       outputTokens: usage.completion_tokens ?? 0,
       reasoningTokens: usage.completion_tokens_details?.reasoning_tokens ?? 0,
       cacheReadTokens: usage.prompt_tokens_details?.cached_tokens ?? 0,
-      //cacheWriteTokens = usage.providerMetadata?.["anthropic"]?.["cacheCreationInputTokens"] ?? 0
-      cacheWriteTokens: 0,
     }),
   })
 }
diff --git a/cloud/app/src/routes/zen/v1/messages.ts b/cloud/app/src/routes/zen/v1/messages.ts
new file mode 100644
index 000000000..ff399b034
--- /dev/null
+++ b/cloud/app/src/routes/zen/v1/messages.ts
@@ -0,0 +1,61 @@
+import type { APIEvent } from "@solidjs/start/server"
+import { handler } from "~/util/zen"
+
+type Usage = {
+  cache_creation?: {
+    ephemeral_5m_input_tokens?: number
+    ephemeral_1h_input_tokens?: number
+  }
+  cache_creation_input_tokens?: number
+  cache_read_input_tokens?: number
+  input_tokens?: number
+  output_tokens?: number
+  server_tool_use?: {
+    web_search_requests?: number
+  }
+}
+
+export function POST(input: APIEvent) {
+  let usage: Usage
+  return handler(input, {
+    modifyBody: (body: any) => ({
+      ...body,
+      service_tier: "standard_only",
+    }),
+    setAuthHeader: (headers: Headers, apiKey: string) => headers.set("x-api-key", apiKey),
+    parseApiKey: (headers: Headers) => headers.get("x-api-key") ?? undefined,
+    onStreamPart: (chunk: string) => {
+      const data = chunk.split("\n")[1]
+      if (!data.startsWith("data: ")) return
+
+      let json
+      try {
+        json = JSON.parse(data.slice(6)) as { usage?: Usage }
+      } catch (e) {
+        return
+      }
+
+      if (!json.usage) return
+      usage = {
+        ...usage,
+        ...json.usage,
+        cache_creation: {
+          ...usage?.cache_creation,
+          ...json.usage.cache_creation,
+        },
+        server_tool_use: {
+          ...usage?.server_tool_use,
+          ...json.usage.server_tool_use,
+        },
+      }
+    },
+    getStreamUsage: () => usage,
+    normalizeUsage: (usage: Usage) => ({
+      inputTokens: usage.input_tokens ?? 0,
+      outputTokens: usage.output_tokens ?? 0,
+      cacheReadTokens: usage.cache_read_input_tokens ?? 0,
+      cacheWrite5mTokens: usage.cache_creation?.ephemeral_5m_input_tokens,
+      cacheWrite1hTokens: usage.cache_creation?.ephemeral_1h_input_tokens,
+    }),
+  })
+}
diff --git a/cloud/app/src/routes/zen/v1/responses.ts b/cloud/app/src/routes/zen/v1/responses.ts
index 844b4fef8..518a431ab 100644
--- a/cloud/app/src/routes/zen/v1/responses.ts
+++ b/cloud/app/src/routes/zen/v1/responses.ts
@@ -2,8 +2,13 @@ import type { APIEvent } from "@solidjs/start/server"
 import { handler } from "~/util/zen"
 
 export function POST(input: APIEvent) {
+  let usage: any
   return handler(input, {
-    parseUsageChunk: (chunk: string) => {
+    setAuthHeader: (headers: Headers, apiKey: string) => {
+      headers.set("authorization", `Bearer ${apiKey}`)
+    },
+    parseApiKey: (headers: Headers) => headers.get("authorization")?.split(" ")[1],
+    onStreamPart: (chunk: string) => {
       const [event, data] = chunk.split("\n")
       if (event !== "event: response.completed") return
       if (!data.startsWith("data: ")) return
@@ -15,9 +20,11 @@ export function POST(input: APIEvent) {
         return
       }
 
-      return json.response?.usage
+      if (!json.response?.usage) return
+      usage = json.response.usage
     },
-    buildUsage: (usage: any) => {
+    getStreamUsage: () => usage,
+    normalizeUsage: (usage: any) => {
       const inputTokens = usage.input_tokens ?? 0
       const outputTokens = usage.output_tokens ?? 0
       const reasoningTokens = usage.output_tokens_details?.reasoning_tokens ?? 0
@@ -27,7 +34,6 @@ export function POST(input: APIEvent) {
         outputTokens: outputTokens - reasoningTokens,
         reasoningTokens,
         cacheReadTokens,
-        cacheWriteTokens: 0,
       }
     },
   })