fix: chat query (#1170)

Co-authored-by: josh <josh@afterima.ge>
2025-09-09 00:29:04 +02:00 · 2025-09-09 00:29:04 +02:00 · f2320bf321
commit f2320bf321
parent d3b3449eb8
14 changed files with 879 additions and 145 deletions
--- a/app/(chat)/api/chat/[id]/stream/route.ts
+++ b/app/(chat)/api/chat/[id]/stream/route.ts
@ -34,7 +34,7 @@ export async function GET(
    return new ChatSDKError('unauthorized:chat').toResponse();
  }

-  let chat: Chat;
+  let chat: Chat | null;

  try {
    chat = await getChatById({ id: chatId });
--- a/app/(chat)/api/chat/route.ts
+++ b/app/(chat)/api/chat/route.ts
@ -18,6 +18,7 @@ import {
  saveChat,
  saveMessages,
 } from '@/lib/db/queries';
+import { updateChatLastContextById } from '@/lib/db/queries';
 import { convertToUIMessages, generateUUID } from '@/lib/utils';
 import { generateTitleFromUserMessage } from '../../actions';
 import { createDocument } from '@/lib/ai/tools/create-document';
@ -208,6 +209,17 @@ export async function POST(request: Request) {
            chatId: id,
          })),
        });
+
+        if (finalUsage) {
+          try {
+            await updateChatLastContextById({
+              chatId: id,
+              context: finalUsage,
+            });
+          } catch (err) {
+            console.warn('Unable to persist last usage for chat', id, err);
+          }
+        }
      },
      onError: () => {
        return 'Oops, an error occurred!';
@ -251,7 +263,7 @@ export async function DELETE(request: Request) {

  const chat = await getChatById({ id });

-  if (chat.userId !== session.user.id) {
+  if (chat?.userId !== session.user.id) {
    return new ChatSDKError('forbidden:chat').toResponse();
  }

--- a/app/(chat)/chat/[id]/page.tsx
+++ b/app/(chat)/chat/[id]/page.tsx
@ -7,6 +7,7 @@ import { getChatById, getMessagesByChatId } from '@/lib/db/queries';
 import { DataStreamHandler } from '@/components/data-stream-handler';
 import { DEFAULT_CHAT_MODEL } from '@/lib/ai/models';
 import { convertToUIMessages } from '@/lib/utils';
+import { LanguageModelV2Usage } from '@ai-sdk/provider';

 export default async function Page(props: { params: Promise<{ id: string }> }) {
  const params = await props.params;
@ -53,6 +54,7 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
          isReadonly={session?.user?.id !== chat.userId}
          session={session}
          autoResume={true}
+          initialLastContext={chat.lastContext ?? undefined}
        />
        <DataStreamHandler />
      </>
@ -69,6 +71,7 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
        isReadonly={session?.user?.id !== chat.userId}
        session={session}
        autoResume={true}
+        initialLastContext={chat.lastContext ?? undefined}
      />
      <DataStreamHandler />
    </>
--- a/components/chat.tsx
+++ b/components/chat.tsx
@ -31,6 +31,7 @@ export function Chat({
  isReadonly,
  session,
  autoResume,
+  initialLastContext,
 }: {
  id: string;
  initialMessages: ChatMessage[];
@ -39,6 +40,7 @@ export function Chat({
  isReadonly: boolean;
  session: Session;
  autoResume: boolean;
+  initialLastContext?: LanguageModelUsage;
 }) {
  const { visibilityType } = useChatVisibility({
    chatId: id,
@ -49,7 +51,9 @@ export function Chat({
  const { setDataStream } = useDataStream();

  const [input, setInput] = useState<string>('');
-  const [usage, setUsage] = useState<LanguageModelUsage | undefined>(undefined);
+  const [usage, setUsage] = useState<LanguageModelUsage | undefined>(
+    initialLastContext,
+  );

  const {
    messages,
--- a/components/elements/context.tsx
+++ b/components/elements/context.tsx
@ -10,6 +10,7 @@ import type { ComponentProps } from 'react';
 import type { LanguageModelUsage } from 'ai';
 import { breakdownTokens, estimateCost, normalizeUsage } from 'tokenlens';
 import { Separator } from '@/components/ui/separator';
+import { Progress } from '@/components/ui/progress';

 export type ContextProps = ComponentProps<'button'> & {
  /** Total context window size in tokens */
@ -20,8 +21,6 @@ export type ContextProps = ComponentProps<'button'> & {
  usage?: LanguageModelUsage | undefined;
  /** Optional model id (canonical or alias) to compute cost */
  modelId?: string;
-  /** Show token breakdown and optional cost inside hover */
-  showBreakdown?: boolean;
 };

 const THOUSAND = 1000;
@ -81,6 +80,11 @@ const formatUSD = (value?: number) => {
  return `$${trimmed}`;
 };

+const formatUSDFixed = (value?: number, decimals = 5) => {
+  if (value === undefined || !Number.isFinite(value)) return undefined;
+  return `$${Number(value).toFixed(decimals)}`;
+};
+
 type ContextIconProps = {
  percent: number; // 0 - 100
 };
@ -125,13 +129,46 @@ export const ContextIcon = ({ percent }: ContextIconProps) => {
  );
 };

+function TokensWithCost({
+  tokens,
+  costText,
+}: {
+  tokens?: number;
+  costText?: string;
+}) {
+  return (
+    <span>
+      {tokens === undefined ? '—' : formatTokens(tokens)}
+      {costText ? (
+        <span className="ml-2 text-muted-foreground">• {costText}</span>
+      ) : null}
+    </span>
+  );
+}
+
+function InfoRow({
+  label,
+  tokens,
+  costText,
+}: {
+  label: string;
+  tokens?: number;
+  costText?: string;
+}) {
+  return (
+    <div className="flex items-center justify-between text-xs">
+      <span className="text-muted-foreground">{label}</span>
+      <TokensWithCost tokens={tokens} costText={costText} />
+    </div>
+  );
+}
+
 export const Context = ({
  className,
  maxTokens,
  usedTokens,
  usage,
  modelId,
-  showBreakdown,
  ...props
 }: ContextProps) => {
  const safeMax = Math.max(0, Number.isFinite(maxTokens) ? maxTokens : 0);
@ -139,38 +176,106 @@ export const Context = ({
    Math.max(0, Number.isFinite(usedTokens) ? usedTokens : 0),
    safeMax,
  );
+
+  // used percent and used tokens to display (demo-aware)
+  const displayUsedTokens = safeUsed;
  const usedPercent =
    safeMax > 0
-      ? Math.min(PERCENT_MAX, Math.max(0, (safeUsed / safeMax) * PERCENT_MAX))
+      ? Math.min(
+          PERCENT_MAX,
+          Math.max(0, (displayUsedTokens / safeMax) * PERCENT_MAX),
+        )
      : 0;

  const displayPct = formatPercent(Math.round(usedPercent * 10) / 10);

-  const used = formatTokens(safeUsed);
+  const used = formatTokens(displayUsedTokens);
  const total = formatTokens(safeMax);

-  const uNorm = normalizeUsage(usage as any);
-  const uBreakdown = breakdownTokens(usage as any);
-  const costUSD = modelId
-    ? estimateCost({ modelId, usage: uNorm }).totalUSD
-    : undefined;
-  const costText = formatUSD(costUSD);
+  const uNorm = normalizeUsage(usage);
+  const uBreakdown = breakdownTokens(usage);
+
+  const hasUsage =
+    !!usage &&
+    ((uNorm.input ?? 0) > 0 ||
+      (uNorm.output ?? 0) > 0 ||
+      (uBreakdown.cacheReads ?? 0) > 0 ||
+      (uBreakdown.cacheWrites ?? 0) > 0 ||
+      (uBreakdown.reasoningTokens ?? 0) > 0);
+
+  // Values to render in rows (demo or real)
+  const displayInput = uNorm.input;
+  const displayOutput = uNorm.output;
+
+  // Per-segment costs
+  const inputCostText = modelId
+    ? formatUSDFixed(
+        estimateCost({
+          modelId,
+          usage: { input: displayInput ?? 0, output: 0 },
+        }).inputUSD,
+      )
+    : undefined;
+  const outputCostText = modelId
+    ? formatUSDFixed(
+        estimateCost({
+          modelId,
+          usage: { input: 0, output: displayOutput ?? 0 },
+        }).outputUSD,
+      )
+    : undefined;
+  // Not supported by tokenlens pricing hints; leave undefined so no bullet is shown
+  const cacheReadsTokens = uBreakdown.cacheReads ?? 0;
+  const cacheWritesTokens = uBreakdown.cacheWrites ?? 0;
+  const cacheReadsCostText =
+    modelId && cacheReadsTokens > 0
+      ? formatUSDFixed(
+          estimateCost({
+            modelId,
+            // Cast to any to support extended pricing fields provided by tokenlens
+            usage: { cacheReads: cacheReadsTokens } as any,
+          }).totalUSD,
+        )
+      : undefined;
+  const cacheWritesCostText =
+    modelId && cacheWritesTokens > 0
+      ? formatUSDFixed(
+          estimateCost({
+            modelId,
+            usage: { cacheWrites: cacheWritesTokens } as any,
+          }).totalUSD,
+        )
+      : undefined;
+
+  const reasoningTokens = uBreakdown.reasoningTokens ?? 0;
+  let reasoningCostText: string | undefined;
+  if (modelId && reasoningTokens > 0) {
+    const est = estimateCost({
+      modelId,
+      usage: { reasoningTokens },
+    }).totalUSD;
+    // TokenLens does not provide reasoning pricing for some models. Show em dash when unknown.
+    reasoningCostText =
+      est && Number.isFinite(est) && est > 0 ? formatUSDFixed(est) : '—';
+  }
+
+  const costUSD = modelId
+    ? estimateCost({
+        modelId,
+        usage: { input: displayInput ?? 0, output: displayOutput ?? 0 },
+      }).totalUSD
+    : undefined;
+  const costText = formatUSDFixed(costUSD);

-  const segInput = Math.max(0, uNorm.input ?? 0);
-  const segOutput = Math.max(0, uNorm.output ?? 0);
-  const segCacheR = Math.max(0, uBreakdown.cacheReads ?? 0);
-  const segCacheW = Math.max(0, uBreakdown.cacheWrites ?? 0);
-  const denom = safeMax > 0 ? safeMax : 1;
-  const w = (n: number) =>
-    `${Math.min(100, Math.max(0, (n / denom) * 100)).toFixed(2)}%`;
  const fmtOrUnknown = (n?: number) =>
    n === undefined ? '—' : formatTokens(n);
+
  return (
    <HoverCard closeDelay={100} openDelay={100}>
      <HoverCardTrigger asChild>
        <button
          className={cn(
-            'inline-flex select-none items-center gap-2 rounded-md px-2.5 py-1 text-sm',
+            'inline-flex select-none items-center gap-1 rounded-md px-2.5 py-1 text-sm',
            'bg-background text-foreground',
          )}
          type="button"
@ -182,111 +287,56 @@ export const Context = ({
          <ContextIcon percent={usedPercent} />
        </button>
      </HoverCardTrigger>
-      <HoverCardContent align="center" className="w-fit p-3">
+      <HoverCardContent align="end" side="top" className="w-fit p-3">
        <div className="min-w-[240px] space-y-2">
-          <p className="text-center text-sm">
+          <p className="text-start text-sm">
            {displayPct} • {used} / {total} tokens
-            {costText ? (
-              <span className="ml-1 text-muted-foreground">• {costText}</span>
-            ) : null}
          </p>
-          {true && (
-            <div className="space-y-2">
-              <div className="h-2 w-full overflow-hidden rounded-full bg-muted">
-                <div
-                  className="h-full"
-                  style={{
-                    width: w(segCacheR),
-                    background: 'var(--chart-2)',
-                    opacity: 0.9,
-                  }}
+          <div className="space-y-2">
+            <Progress className="h-2 bg-muted" value={usedPercent} />
+          </div>
+          <div className="mt-1 space-y-1">
+            {hasUsage && uBreakdown.cacheReads && uBreakdown.cacheReads > 0 && (
+              <InfoRow
+                label="Cache Hits"
+                tokens={uBreakdown.cacheReads}
+                costText={cacheReadsCostText}
+              />
+            )}
+            {hasUsage &&
+              uBreakdown.cacheWrites &&
+              uBreakdown.cacheWrites > 0 && (
+                <InfoRow
+                  label="Cache Writes"
+                  tokens={uBreakdown.cacheWrites}
+                  costText={cacheWritesCostText}
                />
-                <div
-                  className="h-full"
-                  style={{
-                    width: w(segCacheW),
-                    background: 'var(--chart-4)',
-                    opacity: 0.9,
-                  }}
-                />
-                <div
-                  className="h-full"
-                  style={{
-                    width: w(segInput),
-                    background: 'var(--chart-1)',
-                    opacity: 0.9,
-                  }}
-                />
-                <div
-                  className="h-full"
-                  style={{
-                    width: w(segOutput),
-                    background: 'var(--chart-3)',
-                    opacity: 0.9,
-                  }}
-                />
-              </div>
-              <div className="grid grid-cols-2 gap-x-4 gap-y-1 text-xs">
-                <div className="flex items-center justify-between">
-                  <span className="flex items-center gap-2 text-muted-foreground">
-                    <span className="inline-block size-2 rounded-sm bg-chart-1" />
-                    Cache Hits
-                  </span>
-                  <span>{fmtOrUnknown(uBreakdown.cacheReads)}</span>
-                </div>
-                <div className="flex items-center justify-between">
-                  <span className="flex items-center gap-2 text-muted-foreground">
-                    <span className="inline-block size-2 rounded-sm bg-chart-2" />
-                    Cache Writes
-                  </span>
-                  <span>{fmtOrUnknown(uBreakdown.cacheWrites)}</span>
-                </div>
-                <div className="flex items-center justify-between">
-                  <span className="flex items-center gap-2 text-muted-foreground">
-                    <span className="inline-block size-2 rounded-sm bg-chart-3" />
-                    Input
-                  </span>
-                  <span>{formatTokens(uNorm.input)}</span>
-                </div>
-                <div className="flex items-center justify-between">
-                  <span className="flex items-center gap-2 text-muted-foreground">
-                    <span className="inline-block size-2 rounded-sm bg-chart-4" />
-                    Output
-                  </span>
-                  <span>{formatTokens(uNorm.output)}</span>
-                </div>
-              </div>
-            </div>
-          )}
-          {showBreakdown && (
-            <div className="mt-1 space-y-1">
-              <div className="flex items-center justify-between text-xs">
-                <span className="text-muted-foreground">Cache Hits</span>
-                <span>{fmtOrUnknown(uBreakdown.cacheReads)}</span>
-              </div>
-              <div className="flex items-center justify-between text-xs">
-                <span className="text-muted-foreground">Cache Writes</span>
-                <span>{fmtOrUnknown(uBreakdown.cacheWrites)}</span>
-              </div>
-              <div className="flex items-center justify-between text-xs">
-                <span className="text-muted-foreground">Input</span>
-                <span>{formatTokens(uNorm.input)}</span>
-              </div>
-              <div className="flex items-center justify-between text-xs">
-                <span className="text-muted-foreground">Output</span>
-                <span>{formatTokens(uNorm.output)}</span>
-              </div>
-              {costText && (
-                <>
-                  <Separator className="mt-1" />
-                  <div className="flex items-center justify-between pt-1 text-xs">
-                    <span className="text-muted-foreground">Total cost</span>
-                    <span>{costText}</span>
-                  </div>
-                </>
              )}
-            </div>
-          )}
+            <InfoRow
+              label="Input"
+              tokens={displayInput}
+              costText={inputCostText}
+            />
+            <InfoRow
+              label="Output"
+              tokens={displayOutput}
+              costText={outputCostText}
+            />
+            <InfoRow
+              label="Reasoning"
+              tokens={reasoningTokens > 0 ? reasoningTokens : undefined}
+              costText={reasoningCostText}
+            />
+            {costText && (
+              <>
+                <Separator className="mt-1" />
+                <div className="flex items-center justify-between pt-1 text-xs">
+                  <span className="text-muted-foreground">Total cost</span>
+                  <span>{costText}</span>
+                </div>
+              </>
+            )}
+          </div>
        </div>
      </HoverCardContent>
    </HoverCard>
--- a/components/multimodal-input.tsx
+++ b/components/multimodal-input.tsx
@ -214,7 +214,6 @@ function PureMultimodalInput({
      usedTokens,
      usage,
      modelId: modelResolver.modelId,
-      showBreakdown: true as const,
    }),
    [contextMax, usedTokens, usage, modelResolver],
  );
@ -343,20 +342,22 @@ function PureMultimodalInput({
            ))}
          </div>
        )}
-
-        <PromptInputTextarea
-          data-testid="multimodal-input"
-          ref={textareaRef}
-          placeholder="Send a message..."
-          value={input}
-          onChange={handleInput}
-          minHeight={44}
-          maxHeight={200}
-          disableAutoResize={true}
-          className="text-sm resize-none p-3 [&::-webkit-scrollbar]:hidden [-ms-overflow-style:none] [scrollbar-width:none] bg-transparent !border-0 !border-none outline-none ring-0 focus-visible:ring-0 focus-visible:ring-offset-0 focus-visible:outline-none placeholder:text-muted-foreground"
-          rows={1}
-          autoFocus
-        />
+        <div className="flex flex-row gap-2 items-start">
+          <PromptInputTextarea
+            data-testid="multimodal-input"
+            ref={textareaRef}
+            placeholder="Send a message..."
+            value={input}
+            onChange={handleInput}
+            minHeight={44}
+            maxHeight={200}
+            disableAutoResize={true}
+            className="text-sm flex-grow resize-none py-3 px-3 [&::-webkit-scrollbar]:hidden [-ms-overflow-style:none] [scrollbar-width:none] bg-transparent !border-0 !border-none outline-none ring-0 focus-visible:ring-0 focus-visible:ring-offset-0 focus-visible:outline-none placeholder:text-muted-foreground"
+            rows={1}
+            autoFocus
+          />{' '}
+          <Context {...contextProps} />
+        </div>
        <PromptInputToolbar className="px-3 py-2 !border-t-0 !border-top-0 shadow-none dark:!border-transparent dark:border-0">
          <PromptInputTools className="gap-2">
            <AttachmentsButton
@ -365,8 +366,8 @@ function PureMultimodalInput({
              selectedModelId={selectedModelId}
            />
            <ModelSelectorCompact selectedModelId={selectedModelId} />
-            <Context {...contextProps} />
          </PromptInputTools>
+
          {status === 'submitted' ? (
            <StopButton stop={stop} setMessages={setMessages} />
          ) : (
--- a/components/ui/progress.tsx
+++ b/components/ui/progress.tsx
@ -0,0 +1,28 @@
+"use client"
+
+import * as React from "react"
+import * as ProgressPrimitive from "@radix-ui/react-progress"
+
+import { cn } from "@/lib/utils"
+
+const Progress = React.forwardRef<
+  React.ElementRef<typeof ProgressPrimitive.Root>,
+  React.ComponentPropsWithoutRef<typeof ProgressPrimitive.Root>
+>(({ className, value, ...props }, ref) => (
+  <ProgressPrimitive.Root
+    ref={ref}
+    className={cn(
+      "relative h-4 w-full overflow-hidden rounded-full bg-secondary",
+      className
+    )}
+    {...props}
+  >
+    <ProgressPrimitive.Indicator
+      className="h-full w-full flex-1 bg-primary transition-all"
+      style={{ transform: `translateX(-${100 - (value || 0)}%)` }}
+    />
+  </ProgressPrimitive.Root>
+))
+Progress.displayName = ProgressPrimitive.Root.displayName
+
+export { Progress }
--- a/lib/db/migrations/0007_flowery_ben_parker.sql
+++ b/lib/db/migrations/0007_flowery_ben_parker.sql
@ -0,0 +1 @@
+ALTER TABLE "Chat" ADD COLUMN "lastContext" jsonb;
--- a/lib/db/migrations/meta/0007_snapshot.json
+++ b/lib/db/migrations/meta/0007_snapshot.json
@ -0,0 +1,571 @@
+{
+  "id": "097660a7-976a-4b3e-8ebb-79312e3ece6f",
+  "prevId": "443de550-b7e8-4bfb-b229-c12dc6c132f0",
+  "version": "7",
+  "dialect": "postgresql",
+  "tables": {
+    "public.Chat": {
+      "name": "Chat",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "title": {
+          "name": "title",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "visibility": {
+          "name": "visibility",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "'private'"
+        },
+        "lastContext": {
+          "name": "lastContext",
+          "type": "jsonb",
+          "primaryKey": false,
+          "notNull": false
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Chat_userId_User_id_fk": {
+          "name": "Chat_userId_User_id_fk",
+          "tableFrom": "Chat",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Document": {
+      "name": "Document",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "title": {
+          "name": "title",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "content": {
+          "name": "content",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "text": {
+          "name": "text",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "'text'"
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Document_userId_User_id_fk": {
+          "name": "Document_userId_User_id_fk",
+          "tableFrom": "Document",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Document_id_createdAt_pk": {
+          "name": "Document_id_createdAt_pk",
+          "columns": [
+            "id",
+            "createdAt"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.Message_v2": {
+      "name": "Message_v2",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "role": {
+          "name": "role",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "parts": {
+          "name": "parts",
+          "type": "json",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "attachments": {
+          "name": "attachments",
+          "type": "json",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Message_v2_chatId_Chat_id_fk": {
+          "name": "Message_v2_chatId_Chat_id_fk",
+          "tableFrom": "Message_v2",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Message": {
+      "name": "Message",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "role": {
+          "name": "role",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "content": {
+          "name": "content",
+          "type": "json",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Message_chatId_Chat_id_fk": {
+          "name": "Message_chatId_Chat_id_fk",
+          "tableFrom": "Message",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Stream": {
+      "name": "Stream",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Stream_chatId_Chat_id_fk": {
+          "name": "Stream_chatId_Chat_id_fk",
+          "tableFrom": "Stream",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Stream_id_pk": {
+          "name": "Stream_id_pk",
+          "columns": [
+            "id"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.Suggestion": {
+      "name": "Suggestion",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "documentId": {
+          "name": "documentId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "documentCreatedAt": {
+          "name": "documentCreatedAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "originalText": {
+          "name": "originalText",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "suggestedText": {
+          "name": "suggestedText",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "description": {
+          "name": "description",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "isResolved": {
+          "name": "isResolved",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true,
+          "default": false
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Suggestion_userId_User_id_fk": {
+          "name": "Suggestion_userId_User_id_fk",
+          "tableFrom": "Suggestion",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "Suggestion_documentId_documentCreatedAt_Document_id_createdAt_fk": {
+          "name": "Suggestion_documentId_documentCreatedAt_Document_id_createdAt_fk",
+          "tableFrom": "Suggestion",
+          "tableTo": "Document",
+          "columnsFrom": [
+            "documentId",
+            "documentCreatedAt"
+          ],
+          "columnsTo": [
+            "id",
+            "createdAt"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Suggestion_id_pk": {
+          "name": "Suggestion_id_pk",
+          "columns": [
+            "id"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.User": {
+      "name": "User",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "email": {
+          "name": "email",
+          "type": "varchar(64)",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "password": {
+          "name": "password",
+          "type": "varchar(64)",
+          "primaryKey": false,
+          "notNull": false
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Vote_v2": {
+      "name": "Vote_v2",
+      "schema": "",
+      "columns": {
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "messageId": {
+          "name": "messageId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "isUpvoted": {
+          "name": "isUpvoted",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Vote_v2_chatId_Chat_id_fk": {
+          "name": "Vote_v2_chatId_Chat_id_fk",
+          "tableFrom": "Vote_v2",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "Vote_v2_messageId_Message_v2_id_fk": {
+          "name": "Vote_v2_messageId_Message_v2_id_fk",
+          "tableFrom": "Vote_v2",
+          "tableTo": "Message_v2",
+          "columnsFrom": [
+            "messageId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Vote_v2_chatId_messageId_pk": {
+          "name": "Vote_v2_chatId_messageId_pk",
+          "columns": [
+            "chatId",
+            "messageId"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.Vote": {
+      "name": "Vote",
+      "schema": "",
+      "columns": {
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "messageId": {
+          "name": "messageId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "isUpvoted": {
+          "name": "isUpvoted",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Vote_chatId_Chat_id_fk": {
+          "name": "Vote_chatId_Chat_id_fk",
+          "tableFrom": "Vote",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "Vote_messageId_Message_id_fk": {
+          "name": "Vote_messageId_Message_id_fk",
+          "tableFrom": "Vote",
+          "tableTo": "Message",
+          "columnsFrom": [
+            "messageId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Vote_chatId_messageId_pk": {
+          "name": "Vote_chatId_messageId_pk",
+          "columns": [
+            "chatId",
+            "messageId"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    }
+  },
+  "enums": {},
+  "schemas": {},
+  "sequences": {},
+  "_meta": {
+    "columns": {},
+    "schemas": {},
+    "tables": {}
+  }
+}
--- a/lib/db/migrations/meta/_journal.json
+++ b/lib/db/migrations/meta/_journal.json
@ -50,6 +50,13 @@
      "when": 1746118166211,
      "tag": "0006_marvelous_frog_thor",
      "breakpoints": true
+    },
+    {
+      "idx": 7,
+      "version": "7",
+      "when": 1757362773211,
+      "tag": "0007_flowery_ben_parker",
+      "breakpoints": true
    }
  ]
 }
--- a/lib/db/queries.ts
+++ b/lib/db/queries.ts
@ -33,6 +33,7 @@ import { generateUUID } from '../utils';
 import { generateHashedPassword } from './utils';
 import type { VisibilityType } from '@/components/visibility-selector';
 import { ChatSDKError } from '../errors';
+import { LanguageModelV2Usage } from '@ai-sdk/provider';

 // Optionally, if not using email/pass login, you can
 // use the Drizzle adapter for Auth.js / NextAuth
@ -202,6 +203,10 @@ export async function getChatsByUserId({
 export async function getChatById({ id }: { id: string }) {
  try {
    const [selectedChat] = await db.select().from(chat).where(eq(chat.id, id));
+    if (!selectedChat) {
+      return null;
+    }
+
    return selectedChat;
  } catch (error) {
    throw new ChatSDKError('bad_request:database', 'Failed to get chat by id');
@ -469,10 +474,32 @@ export async function updateChatVisiblityById({
  }
 }

+export async function updateChatLastContextById({
+  chatId,
+  context,
+}: {
+  chatId: string;
+  // Store raw LanguageModelUsage to keep it simple
+  context: LanguageModelV2Usage;
+}) {
+  try {
+    return await db
+      .update(chat)
+      .set({ lastContext: context })
+      .where(eq(chat.id, chatId));
+  } catch (error) {
+    console.warn('Failed to update lastContext for chat', chatId, error);
+    return;
+  }
+}
+
 export async function getMessageCountByUserId({
  id,
  differenceInHours,
-}: { id: string; differenceInHours: number }) {
+}: {
+  id: string;
+  differenceInHours: number;
+}) {
  try {
    const twentyFourHoursAgo = new Date(
      Date.now() - differenceInHours * 60 * 60 * 1000,
--- a/lib/db/schema.ts
+++ b/lib/db/schema.ts
@ -4,12 +4,14 @@ import {
  varchar,
  timestamp,
  json,
+  jsonb,
  uuid,
  text,
  primaryKey,
  foreignKey,
  boolean,
 } from 'drizzle-orm/pg-core';
+import type { LanguageModelV2Usage } from '@ai-sdk/provider';

 export const user = pgTable('User', {
  id: uuid('id').primaryKey().notNull().defaultRandom(),
@ -29,6 +31,7 @@ export const chat = pgTable('Chat', {
  visibility: varchar('visibility', { enum: ['public', 'private'] })
    .notNull()
    .default('private'),
+  lastContext: jsonb('lastContext').$type<LanguageModelV2Usage | null>(),
 });

 export type Chat = InferSelectModel<typeof chat>;
--- a/package.json
+++ b/package.json
@ -39,6 +39,7 @@
    "@radix-ui/react-hover-card": "^1.1.15",
    "@radix-ui/react-icons": "^1.3.0",
    "@radix-ui/react-label": "^2.1.0",
+    "@radix-ui/react-progress": "^1.1.7",
    "@radix-ui/react-scroll-area": "^1.2.10",
    "@radix-ui/react-select": "^2.1.2",
    "@radix-ui/react-separator": "^1.1.0",
@ -95,7 +96,7 @@
    "swr": "^2.2.5",
    "tailwind-merge": "^2.5.2",
    "tailwindcss-animate": "^1.0.7",
-    "tokenlens": "^1.0.0",
+    "tokenlens": "^1.1.2",
    "use-stick-to-bottom": "^1.1.1",
    "usehooks-ts": "^3.1.0",
    "zod": "^3.25.76"
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@ -68,6 +68,9 @@ importers:
      '@radix-ui/react-label':
        specifier: ^2.1.0
        version: 2.1.2(@types/react-dom@18.3.5(@types/react@18.3.18))(@types/react@18.3.18)(react-dom@19.0.0-rc-45804af1-20241021(react@19.0.0-rc-45804af1-20241021))(react@19.0.0-rc-45804af1-20241021)
+      '@radix-ui/react-progress':
+        specifier: ^1.1.7
+        version: 1.1.7(@types/react-dom@18.3.5(@types/react@18.3.18))(@types/react@18.3.18)(react-dom@19.0.0-rc-45804af1-20241021(react@19.0.0-rc-45804af1-20241021))(react@19.0.0-rc-45804af1-20241021)
      '@radix-ui/react-scroll-area':
        specifier: ^1.2.10
        version: 1.2.10(@types/react-dom@18.3.5(@types/react@18.3.18))(@types/react@18.3.18)(react-dom@19.0.0-rc-45804af1-20241021(react@19.0.0-rc-45804af1-20241021))(react@19.0.0-rc-45804af1-20241021)
@ -237,8 +240,8 @@ importers:
        specifier: ^1.0.7
        version: 1.0.7(tailwindcss@3.4.17)
      tokenlens:
-        specifier: ^1.0.0
-        version: 1.0.0
+        specifier: ^1.1.2
+        version: 1.1.2
      use-stick-to-bottom:
        specifier: ^1.1.1
        version: 1.1.1(react@19.0.0-rc-45804af1-20241021)
@ -1627,6 +1630,19 @@ packages:
      '@types/react-dom':
        optional: true

+  '@radix-ui/react-progress@1.1.7':
+    resolution: {integrity: sha512-vPdg/tF6YC/ynuBIJlk1mm7Le0VgW6ub6J2UWnTQ7/D23KXcPI1qy+0vBkgKgd38RCMJavBXpB83HPNFMTb0Fg==}
+    peerDependencies:
+      '@types/react': '*'
+      '@types/react-dom': '*'
+      react: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc
+      react-dom: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc
+    peerDependenciesMeta:
+      '@types/react':
+        optional: true
+      '@types/react-dom':
+        optional: true
+
  '@radix-ui/react-roving-focus@1.1.2':
    resolution: {integrity: sha512-zgMQWkNO169GtGqRvYrzb0Zf8NhMHS2DuEB/TiEmVnpr5OqPU3i8lfbxaAmC2J/KYuIQxyoQQ6DxepyXp61/xw==}
    peerDependencies:
@ -4691,8 +4707,8 @@ packages:
    resolution: {integrity: sha512-65P7iz6X5yEr1cwcgvQxbbIw7Uk3gOy5dIdtZ4rDveLqhrdJP+Li/Hx6tyK0NEb+2GCyneCMJiGqrADCSNk8sQ==}
    engines: {node: '>=8.0'}

-  tokenlens@1.0.0:
-    resolution: {integrity: sha512-IslxKcdsPQv4yTUgzeAe9UoEuvXUW6BZyKvKc4N9D7IWnyUBvSTJ3hq8Wb13fM8CkgaS/0k28vpJIQiVQJrpVA==}
+  tokenlens@1.1.2:
+    resolution: {integrity: sha512-PICXahAiVS/2nyc30SijuYZj5wIm2G9+2U7fdvQbl5pa/hIUpQ8H+RdYOzOrSvva6GMqPM5+6cWd9Hqi81Yv3g==}

  trim-lines@3.0.1:
    resolution: {integrity: sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg==}
@ -6013,6 +6029,16 @@ snapshots:
      '@types/react': 18.3.18
      '@types/react-dom': 18.3.5(@types/react@18.3.18)

+  '@radix-ui/react-progress@1.1.7(@types/react-dom@18.3.5(@types/react@18.3.18))(@types/react@18.3.18)(react-dom@19.0.0-rc-45804af1-20241021(react@19.0.0-rc-45804af1-20241021))(react@19.0.0-rc-45804af1-20241021)':
+    dependencies:
+      '@radix-ui/react-context': 1.1.2(@types/react@18.3.18)(react@19.0.0-rc-45804af1-20241021)
+      '@radix-ui/react-primitive': 2.1.3(@types/react-dom@18.3.5(@types/react@18.3.18))(@types/react@18.3.18)(react-dom@19.0.0-rc-45804af1-20241021(react@19.0.0-rc-45804af1-20241021))(react@19.0.0-rc-45804af1-20241021)
+      react: 19.0.0-rc-45804af1-20241021
+      react-dom: 19.0.0-rc-45804af1-20241021(react@19.0.0-rc-45804af1-20241021)
+    optionalDependencies:
+      '@types/react': 18.3.18
+      '@types/react-dom': 18.3.5(@types/react@18.3.18)
+
  '@radix-ui/react-roving-focus@1.1.2(@types/react-dom@18.3.5(@types/react@18.3.18))(@types/react@18.3.18)(react-dom@19.0.0-rc-45804af1-20241021(react@19.0.0-rc-45804af1-20241021))(react@19.0.0-rc-45804af1-20241021)':
    dependencies:
      '@radix-ui/primitive': 1.1.1
@ -9712,7 +9738,7 @@ snapshots:
    dependencies:
      is-number: 7.0.0

-  tokenlens@1.0.0: {}
+  tokenlens@1.1.2: {}

  trim-lines@3.0.1: {}
				`@ -0,0 +1 @@`
				`ALTER TABLE "Chat" ADD COLUMN "lastContext" jsonb;`