sync

3 files changed, 79 insertions, 28 deletions

diff --git a/app/packages/web/src/components/Share.tsx b/app/packages/web/src/components/Share.tsx
index 292921ec7..78f901f67 100644
--- a/app/packages/web/src/components/Share.tsx
+++ b/app/packages/web/src/components/Share.tsx
@@ -23,6 +23,16 @@ type SessionMessage = UIMessage<{
     created: number
     completed?: number
   }
+  assistant?: {
+    modelID: string;
+    providerID: string;
+    cost: number;
+    tokens: {
+      input: number;
+      output: number;
+      reasoning: number;
+    };
+  };
   sessionID: string
   tool: Record<string, {
     properties: Record<string, any>
@@ -36,11 +46,6 @@ type SessionMessage = UIMessage<{
 type SessionInfo = {
   title: string
   cost?: number
-  tokens?: {
-    input?: number
-    output?: number
-    reasoning?: number
-  }
 }
 
 function getPartTitle(role: string, type: string): string | undefined {
@@ -229,6 +234,26 @@ export default function Share(props: { api: string }) {
     )
   }
 
+  const metrics = createMemo(() => {
+    const result = {
+      cost: 0,
+      tokens: {
+        input: 0,
+        output: 0,
+        reasoning: 0,
+      }
+    }
+    for (const msg of messages()) {
+      const assistant = msg.metadata?.assistant
+      if (!assistant) continue
+      result.cost += assistant.cost
+      result.tokens.input += assistant.tokens.input
+      result.tokens.output += assistant.tokens.output
+      result.tokens.reasoning += assistant.tokens.reasoning
+    }
+    return result
+  })
+
   return (
     <main class={`${styles.root} not-content`}>
       <div class={styles.header}>
@@ -242,25 +267,33 @@ export default function Share(props: { api: string }) {
         <div data-section="row">
           <ul data-section="stats">
             <li>
+              <span data-element-label>Cost</span>
+              {metrics().cost ?
+                <span>{metrics().cost}</span>
+                :
+                <span data-placeholder>&mdash;</span>
+              }
+            </li>
+            <li>
               <span data-element-label>Input Tokens</span>
-              {store.info?.tokens?.input ?
-                <span>{store.info?.tokens?.input}</span>
+              {metrics().tokens.input ?
+                <span>{metrics().tokens.input}</span>
                 :
                 <span data-placeholder>&mdash;</span>
               }
             </li>
             <li>
               <span data-element-label>Output Tokens</span>
-              {store.info?.tokens?.output ?
-                <span>{store.info?.tokens?.output}</span>
+              {metrics().tokens.output ?
+                <span>{metrics().tokens.output}</span>
                 :
                 <span data-placeholder>&mdash;</span>
               }
             </li>
             <li>
               <span data-element-label>Reasoning Tokens</span>
-              {store.info?.tokens?.reasoning ?
-                <span>{store.info?.tokens?.reasoning}</span>
+              {metrics().tokens.reasoning ?
+                <span>{metrics().tokens.reasoning}</span>
                 :
                 <span data-placeholder>&mdash;</span>
               }
diff --git a/js/src/server/server.ts b/js/src/server/server.ts
index 11262f4fd..5fb0dfe62 100644
--- a/js/src/server/server.ts
+++ b/js/src/server/server.ts
@@ -18,6 +18,7 @@ const SessionInfo = Session.Info.openapi({
 const ProviderInfo = Config.Provider.openapi({
   ref: "Provider.Info",
 });
+
 type ProviderInfo = z.output<typeof ProviderInfo>;
 
 export namespace Server {
@@ -210,7 +211,7 @@ export namespace Server {
         ),
         async (c) => {
           const body = c.req.valid("json");
-          const msg = await Session.chat(body as any);
+          const msg = await Session.chat(body);
           return c.json(msg);
         },
       )
diff --git a/js/src/session/session.ts b/js/src/session/session.ts
index 011519048..197d4d398 100644
--- a/js/src/session/session.ts
+++ b/js/src/session/session.ts
@@ -6,7 +6,6 @@ import { Storage } from "../storage/storage";
 import { Log } from "../util/log";
 import {
   convertToModelMessages,
-  generateText,
   stepCountIs,
   streamText,
   type TextUIPart,
@@ -20,7 +19,6 @@ import * as tools from "../tool";
 import { Decimal } from "decimal.js";
 
 import PROMPT_ANTHROPIC from "./prompt/anthropic.txt";
-import PROMPT_TITLE from "./prompt/title.txt";
 
 import type { Tool } from "../tool/tool";
 import { Share } from "../share/share";
@@ -42,6 +40,16 @@ export namespace Session {
   export type Info = z.output<typeof Info>;
 
   export type Message = UIMessage<{
+    assistant?: {
+      modelID: string;
+      providerID: string;
+      cost: number;
+      tokens: {
+        input: number;
+        output: number;
+        reasoning: number;
+      };
+    };
     time: {
       created: number;
       completed?: number;
@@ -71,8 +79,8 @@ export namespace Session {
       },
     };
     log.info("created", result);
-    await Storage.writeJSON("session/info/" + result.id, result);
     state().sessions.set(result.id, result);
+    await Storage.writeJSON("session/info/" + result.id, result);
     return result;
   }
 
@@ -184,6 +192,7 @@ export namespace Session {
       }
       msgs.push(system);
       state().messages.set(input.sessionID, msgs);
+      /*
       generateText({
         messages: convertToModelMessages([
           {
@@ -206,6 +215,7 @@ export namespace Session {
           draft.title = result.text;
         });
       });
+      */
       await write(system);
     }
     const msg: Message = {
@@ -228,6 +238,16 @@ export namespace Session {
       role: "assistant",
       parts: [],
       metadata: {
+        assistant: {
+          cost: 0,
+          tokens: {
+            input: 0,
+            output: 0,
+            reasoning: 0,
+          },
+          modelID: input.modelID,
+          providerID: input.providerID,
+        },
         time: {
           created: Date.now(),
         },
@@ -238,19 +258,16 @@ export namespace Session {
     const controller = new AbortController();
     pending.set(input.sessionID, controller);
     const result = streamText({
-      onStepFinish: (step) => {
-        update(input.sessionID, (draft) => {
-          const input = step.usage.inputTokens ?? 0;
-          const output = step.usage.outputTokens ?? 0;
-          const reasoning = step.usage.reasoningTokens ?? 0;
-          draft.tokens.input += input;
-          draft.tokens.output += output;
-          draft.tokens.reasoning += reasoning;
-          draft.cost = new Decimal(draft.cost ?? 0)
-            .add(new Decimal(input).mul(model.info.cost.input))
-            .add(new Decimal(output).mul(model.info.cost.output))
-            .toNumber();
-        });
+      onStepFinish: async (step) => {
+        const assistant = next.metadata!.assistant!;
+        assistant.tokens.input = step.usage.inputTokens ?? 0;
+        assistant.tokens.output = step.usage.outputTokens ?? 0;
+        assistant.tokens.reasoning = step.usage.reasoningTokens ?? 0;
+        assistant.cost = new Decimal(0)
+          .add(new Decimal(assistant.tokens.input).mul(model.info.cost.input))
+          .add(new Decimal(assistant.tokens.output).mul(model.info.cost.output))
+          .toNumber();
+        await write(next);
       },
       abortSignal: controller.signal,
       maxRetries: 6,