@gram-ai/elements 1.27.0 → 1.27.2

package/package.json

@@ -2,7 +2,7 @@
   "name": "@gram-ai/elements",
   "description": "Gram Elements is a library of UI primitives for building chat-like experiences for MCP Servers.",
   "type": "module",
-  "version": "1.27.0",
+  "version": "1.27.2",
   "main": "dist/index.js",
   "exports": {
     ".": {
@@ -15,6 +15,11 @@
       "import": "./dist/server.js",
       "require": "./dist/server.cjs"
     },
+    "./server/core": {
+      "types": "./dist/server/core.d.ts",
+      "import": "./dist/server/core.js",
+      "require": "./dist/server/core.cjs"
+    },
     "./server/express": {
       "types": "./dist/server/express.d.ts",
       "import": "./dist/server/express.js",

package/src/contexts/ElementsProvider.tsx

@@ -25,8 +25,8 @@ import { Plugin } from '@/types/plugins'
 import {
   AssistantRuntimeProvider,
   AssistantTool,
-  unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime,
   useAssistantState,
+  unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime,
 } from '@assistant-ui/react'
 import {
   frontendTools as convertFrontendToolsToAISDKTools,
@@ -36,6 +36,7 @@ import { createOpenRouter } from '@openrouter/ai-sdk-provider'
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import {
   convertToModelMessages,
+  createUIMessageStream,
   smoothStream,
   stepCountIs,
   streamText,
@@ -52,14 +53,14 @@ import {
   useState,
 } from 'react'
 import { useAuth } from '../hooks/useAuth'
-import { ElementsContext } from './contexts'
-import { ToolApprovalProvider } from './ToolApprovalContext'
+import { ChatIdContext } from './ChatIdContext'
 import {
   ConnectionStatusProvider,
   useConnectionStatusOptional,
 } from './ConnectionStatusContext'
+import { ElementsContext } from './contexts'
+import { ToolApprovalProvider } from './ToolApprovalContext'
 import { ToolExecutionProvider } from './ToolExecutionContext'
-import { ChatIdContext } from './ChatIdContext'
 
 /**
  * Extracts executable tools from frontend tool definitions.
@@ -158,6 +159,10 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
     auth: config.api,
     projectSlug: config.projectSlug,
   })
+
+  // Ref to access ensureValidHeaders in async transport without stale closures
+  const ensureValidHeadersRef = useRef(auth.ensureValidHeaders)
+  ensureValidHeadersRef.current = auth.ensureValidHeaders
   const toolApproval = useToolApproval()
 
   const [model, setModel] = useState<Model>(
@@ -266,6 +271,9 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
           throw new Error('Session is loading')
         }
 
+        // Ensure the session token is still valid; refresh if expired
+        const validHeaders = await ensureValidHeadersRef.current()
+
         // Get chat ID - use the synced remoteId ref first (history mode),
         // fall back to generated ID (non-history mode)
         let chatId = currentRemoteIdRef.current
@@ -318,7 +326,7 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
 
         // Include Gram-Chat-ID header for chat persistence and Gram-Environment for environment selection
         const headersWithChatId = {
-          ...auth.headers,
+          ...validHeaders,
           'Gram-Chat-ID': chatId,
           'X-Gram-Source': 'elements',
           ...config.api?.headers, // We do this after X-Gram-Source so the playground can override it
@@ -327,6 +335,13 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
           }),
         }
 
+        // Update MCP headers with the (possibly refreshed) session token
+        // so mid-stream MCP tool calls use the fresh token
+        const freshSession = validHeaders['Gram-Chat-Session']
+        if (freshSession) {
+          mcpHeaders['Gram-Chat-Session'] = freshSession
+        }
+
         // Create OpenRouter model (only needed when not using custom model)
         const openRouterModel = usingCustomModel
           ? null
@@ -395,7 +410,12 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
           // Mark as connected when stream starts successfully
           connectionStatus?.markConnected()
 
-          return result.toUIMessageStream()
+          // This weird construction is necessary to get errors to propagate properly to assistant-ui
+          return createUIMessageStream({
+            execute: ({ writer }) => {
+              writer.merge(result.toUIMessageStream())
+            },
+          })
         } catch (error) {
           console.error('Error creating stream:', error)
           trackError(error, { source: 'stream-creation' })
@@ -430,7 +450,6 @@ const ElementsProviderInner = ({ children, config }: ElementsProviderProps) => {
       mcpTools,
       getApprovalHelpers,
       apiUrl,
-      auth.headers,
       auth.isLoading,
       connectionStatus,
     ]

package/src/hooks/useAuth.ts

@@ -1,14 +1,17 @@
 import { useReplayContext } from '@/contexts/ReplayContext'
 import {
   hasExplicitSessionAuth,
+  isAnyStaticSession,
   isDangerousApiKeyAuth,
   isStaticSessionAuth,
   isUnifiedFunctionSession,
   isUnifiedStaticSession,
 } from '@/lib/auth'
-import { useMemo } from 'react'
+import { getTokenExpiry } from '@/lib/token'
+import { useCallback, useMemo } from 'react'
 import { ApiConfig, GetSessionFn } from '../types'
-import { useSession } from './useSession'
+import { getChatSessionQueryKey, useSession } from './useSession'
+import { useQueryClient } from '@tanstack/react-query'
 
 declare const __GRAM_API_URL__: string | undefined
 
@@ -16,10 +19,12 @@ export type Auth =
   | {
       headers: Record<string, string>
       isLoading: false
+      ensureValidHeaders: () => Promise<Record<string, string>>
     }
   | {
       headers?: Record<string, string>
       isLoading: true
+      ensureValidHeaders: () => Promise<Record<string, string>>
     }
 
 async function defaultGetSession(init: {
@@ -62,7 +67,7 @@ function createDangerousApiKeySessionFn(
 
 /**
  * Hook to fetch or retrieve the session token for the chat.
- * @returns The session token string or null
+ * @returns Auth object with headers and ensureValidHeaders for pre-request token refresh
  */
 export const useAuth = ({
   projectSlug,
@@ -73,6 +78,7 @@ export const useAuth = ({
 }): Auth => {
   const replayCtx = useReplayContext()
   const isReplay = replayCtx?.isReplay ?? false
+  const queryClient = useQueryClient()
 
   const apiUrl = useMemo(() => {
     const envUrl =
@@ -117,17 +123,58 @@ export const useAuth = ({
     projectSlug,
   })
 
+  const shouldRefresh = !isAnyStaticSession(auth) && !isReplay
+
+  const ensureValidHeaders = useCallback(async (): Promise<
+    Record<string, string>
+  > => {
+    const queryKey = getChatSessionQueryKey(projectSlug)
+    const cachedToken = queryClient.getQueryData<string>(queryKey)
+
+    if (!shouldRefresh || !getSession) {
+      return {
+        'Gram-Project': projectSlug,
+        ...(cachedToken && { 'Gram-Chat-Session': cachedToken }),
+      }
+    }
+
+    // Check if the cached token is expired (or within 30s of expiry).
+    // staleTime=0 forces a refetch; Infinity keeps the cached value.
+    const exp = cachedToken ? getTokenExpiry(cachedToken) : null
+    const isExpired = exp !== null && Date.now() >= exp * 1000 - 30_000
+    const staleTime = isExpired ? 0 : Infinity
+
+    try {
+      const token = await queryClient.fetchQuery({
+        queryKey,
+        queryFn: () => getSession({ projectSlug }),
+        staleTime,
+      })
+      return {
+        'Gram-Project': projectSlug,
+        ...(token && { 'Gram-Chat-Session': token }),
+      }
+    } catch {
+      return {
+        'Gram-Project': projectSlug,
+        ...(cachedToken && { 'Gram-Chat-Session': cachedToken }),
+      }
+    }
+  }, [shouldRefresh, getSession, projectSlug, queryClient])
+
   // In replay mode, return immediately without waiting for session
   if (isReplay) {
     return {
       headers: {},
       isLoading: false,
+      ensureValidHeaders: async () => ({}),
     }
   }
 
   return !session
     ? {
         isLoading: true,
+        ensureValidHeaders,
       }
     : {
         headers: {
@@ -135,5 +182,6 @@ export const useAuth = ({
           'Gram-Chat-Session': session,
         },
         isLoading: false,
+        ensureValidHeaders,
       }
 }

package/src/hooks/useSession.ts

@@ -1,6 +1,10 @@
 import { GetSessionFn } from '@/types'
 import { useQuery, useQueryClient } from '@tanstack/react-query'
 
+export function getChatSessionQueryKey(projectSlug: string) {
+  return ['chatSession', projectSlug] as const
+}
+
 /**
  * Hook to fetch or retrieve the session token for the chat.
  * @returns The session token string or null
@@ -13,21 +17,14 @@ export const useSession = ({
   projectSlug: string
 }): string | null => {
   const queryClient = useQueryClient()
-  const queryKey = ['chatSession', projectSlug]
+  const queryKey = getChatSessionQueryKey(projectSlug)
 
-  const queryState = queryClient.getQueryState(queryKey)
-  const hasData = queryState?.data !== undefined
-  // Check if data is stale - with staleTime: Infinity, data never becomes stale
-  // but we check dataUpdatedAt to determine if we should refetch
-  const dataUpdatedAt = queryState?.dataUpdatedAt ?? 0
-  const staleTime = Infinity // Matches the staleTime in useQuery options
-  const isStale = hasData && Date.now() - dataUpdatedAt > staleTime
-  const shouldFetch = !hasData || isStale
+  const hasData = queryClient.getQueryState(queryKey)?.data !== undefined
 
   const { data: fetchedSessionToken } = useQuery({
     queryKey,
     queryFn: () => getSession!({ projectSlug }),
-    enabled: shouldFetch && getSession !== null,
+    enabled: !hasData && getSession !== null,
     staleTime: Infinity, // Session tokens don't need to be refetched
     gcTime: Infinity, // Keep in cache indefinitely
   })

package/src/lib/token.test.ts

@@ -0,0 +1,79 @@
+import { describe, expect, it } from 'vitest'
+import { getTokenExpiry, isTokenExpired } from './token'
+
+/** Helper: build a JWT with a given payload (no signature verification needed). */
+function makeJwt(payload: Record<string, unknown>): string {
+  const header = btoa(JSON.stringify({ alg: 'HS256', typ: 'JWT' }))
+  const body = btoa(JSON.stringify(payload))
+  return `${header}.${body}.fake-signature`
+}
+
+describe('getTokenExpiry', () => {
+  it('returns exp for a valid JWT', () => {
+    const exp = 1700000000
+    expect(getTokenExpiry(makeJwt({ exp }))).toBe(exp)
+  })
+
+  it('returns null when exp is missing', () => {
+    expect(getTokenExpiry(makeJwt({ sub: 'user' }))).toBeNull()
+  })
+
+  it('returns null for a non-JWT string', () => {
+    expect(getTokenExpiry('not-a-jwt')).toBeNull()
+  })
+
+  it('returns null for an empty string', () => {
+    expect(getTokenExpiry('')).toBeNull()
+  })
+
+  it('returns null when payload is not valid JSON', () => {
+    // Two dots but the middle segment is not valid base64 JSON
+    expect(getTokenExpiry('a.!!!.b')).toBeNull()
+  })
+
+  it('handles base64url characters (- and _)', () => {
+    // Manually craft a payload with base64url-specific chars
+    const payload = { exp: 1700000000 }
+    const json = JSON.stringify(payload)
+    const b64 = btoa(json)
+      .replace(/\+/g, '-')
+      .replace(/\//g, '_')
+      .replace(/=+$/, '')
+    const token = `header.${b64}.sig`
+    expect(getTokenExpiry(token)).toBe(1700000000)
+  })
+})
+
+describe('isTokenExpired', () => {
+  it('returns true for an expired token', () => {
+    // exp = 1 second ago
+    const exp = Math.floor(Date.now() / 1000) - 1
+    expect(isTokenExpired(makeJwt({ exp }))).toBe(true)
+  })
+
+  it('returns true when token is within the buffer window', () => {
+    // exp = 20 seconds from now (within 30s default buffer)
+    const exp = Math.floor(Date.now() / 1000) + 20
+    expect(isTokenExpired(makeJwt({ exp }))).toBe(true)
+  })
+
+  it('returns false when token is well outside the buffer', () => {
+    // exp = 5 minutes from now
+    const exp = Math.floor(Date.now() / 1000) + 300
+    expect(isTokenExpired(makeJwt({ exp }))).toBe(false)
+  })
+
+  it('respects a custom buffer', () => {
+    // exp = 20 seconds from now, buffer = 10s → should NOT be expired
+    const exp = Math.floor(Date.now() / 1000) + 20
+    expect(isTokenExpired(makeJwt({ exp }), 10_000)).toBe(false)
+  })
+
+  it('returns false (fail-open) for a non-JWT string', () => {
+    expect(isTokenExpired('opaque-session-token')).toBe(false)
+  })
+
+  it('returns false (fail-open) for an empty string', () => {
+    expect(isTokenExpired('')).toBe(false)
+  })
+})

package/src/lib/token.ts

@@ -0,0 +1,39 @@
+/**
+ * Extracts the `exp` claim from a JWT token without verifying the signature.
+ * Returns the expiry as a Unix timestamp (seconds), or null if the token
+ * is not a valid JWT or has no `exp` claim.
+ */
+export function getTokenExpiry(token: string): number | null {
+  try {
+    const parts = token.split('.')
+    if (parts.length !== 3) return null
+
+    // base64url → base64 → decode
+    let payload = parts[1].replace(/-/g, '+').replace(/_/g, '/')
+    while (payload.length % 4) payload += '='
+
+    const json = atob(payload)
+    const parsed = JSON.parse(json)
+
+    if (typeof parsed.exp === 'number') {
+      return parsed.exp
+    }
+    return null
+  } catch {
+    return null
+  }
+}
+
+/**
+ * Returns true when the token is expired or within `bufferMs` milliseconds
+ * of expiry. Fails open (returns false) for non-JWT tokens or tokens
+ * without an `exp` claim so they pass through unchanged.
+ */
+export function isTokenExpired(
+  token: string,
+  bufferMs: number = 30_000
+): boolean {
+  const exp = getTokenExpiry(token)
+  if (exp === null) return false // fail-open for non-JWT tokens
+  return Date.now() >= exp * 1000 - bufferMs
+}

package/src/server/tanstack-start.ts

@@ -1,70 +1,37 @@
 /**
  * TanStack Start adapter for Gram Elements server handlers.
  *
- * Provides two approaches for session creation:
+ * Use `createTanStackStartHandler` to create a handler for TanStack Start
+ * server routes that manages chat session creation.
  *
- * 1. **Server Function (RPC)** — use `createTanStackStartSessionFn` to create
- *    a `createServerFn` that can be called directly from client code and passed
- *    to `session` in the Elements config.
- *
- * 2. **API Route** — use `createTanStackStartHandler` to create a handler for
- *    TanStack Start server routes (similar to the Next.js adapter).
- *
- * @example Server Function approach
+ * @example
  * ```typescript
- * // session.functions.ts
- * import { createTanStackStartSessionFn } from '@gram-ai/elements/server/tanstack-start'
+ * // routes/api/chat.session.ts
+ * import { createTanStackStartHandler } from '@gram-ai/elements/server/tanstack-start'
  *
- * export const getSession = createTanStackStartSessionFn({
+ * export const POST = createTanStackStartHandler({
  *   embedOrigin: 'http://localhost:3000',
  *   userIdentifier: 'user-123',
  *   expiresAfter: 3600,
  * })
  * ```
  *
- * @example API Route approach
+ * @example Dynamic configuration
  * ```typescript
  * // routes/api/chat.session.ts
- * import { createTanStackStartHandler } from '@gram-ai/elements/server/tanstack-start'
- *
- * const handler = createTanStackStartHandler({
- *   embedOrigin: 'http://localhost:3000',
- *   userIdentifier: 'user-123',
- *   expiresAfter: 3600,
+ * export const POST = createTanStackStartHandler(async (request) => {
+ *   const user = await getUserFromRequest(request)
+ *   return {
+ *     embedOrigin: 'http://localhost:3000',
+ *     userIdentifier: user.id,
+ *     expiresAfter: 3600,
+ *   }
  * })
  * ```
  */
 
-import { createServerFn } from '@tanstack/react-start'
 import { createChatSession, type SessionHandlerOptions } from './core'
 
-/**
- * Create a TanStack Start server function for session creation.
- *
- * The returned function can be called from client code (RPC-style) and passed
- * to `session` in the Gram Elements config.
- *
- * @param options - Session configuration options
- * @returns A `createServerFn` instance callable from the client
- */
-export function createTanStackStartSessionFn(options: SessionHandlerOptions) {
-  return createServerFn({ method: 'POST' })
-    .inputValidator((data: { projectSlug: string }) => data)
-    .handler(async ({ data }) => {
-      const result = await createChatSession({
-        projectSlug: data.projectSlug,
-        options,
-      })
-
-      if (result.status !== 200) {
-        throw new Error(`Failed to create chat session: ${result.body}`)
-      }
-
-      const parsed = JSON.parse(result.body) as { client_token: string }
-      return parsed.client_token
-    })
-}
-
 /**
  * Create a TanStack Start server route handler for the chat session endpoint.
  *

package/dist/core-Cqad6-xW.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"core-Cqad6-xW.js","sources":["../src/server/core.ts"],"sourcesContent":null,"names":["createChatSession","request","base","response","body","error","errorMessage"],"mappings":"AA2CA,eAAsBA,EACpBC,GACgC;AAChC,QAAMC,IAAO,QAAQ,IAAI,gBAAgB;AAEzC,MAAI;AACF,UAAMC,IAAW,MAAM,MAAMD,IAAO,4BAA4B;AAAA,MAC9D,QAAQ;AAAA,MACR,MAAM,KAAK,UAAU;AAAA,QACnB,cAAcD,EAAQ,QAAQ;AAAA,QAC9B,iBAAiBA,EAAQ,QAAQ;AAAA,QACjC,eAAeA,EAAQ,QAAQ;AAAA,MAAA,CAChC;AAAA,MACD,SAAS;AAAA,QACP,gBAAgB;AAAA,QAChB,gBAAgBA,EAAQ;AAAA,QACxB,YAAYA,EAAQ,QAAQ,UAAU,QAAQ,IAAI,gBAAgB;AAAA,MAAA;AAAA,IACpE,CACD,GAEKG,IAAO,MAAMD,EAAS,KAAA;AAE5B,WAAO;AAAA,MACL,QAAQA,EAAS;AAAA,MACjB,MAAAC;AAAA,MACA,SAAS,EAAE,gBAAgB,mBAAA;AAAA,IAAmB;AAAA,EAElD,SAASC,GAAO;AACd,UAAMC,IACJD,aAAiB,QAAQA,EAAM,UAAU;AAC3C,mBAAQ,MAAM,kCAAkCA,CAAK,GAE9C;AAAA,MACL,QAAQ;AAAA,MACR,MAAM,KAAK,UAAU;AAAA,QACnB,OAAO,oCAAoCC;AAAA,MAAA,CAC5C;AAAA,MACD,SAAS,EAAE,gBAAgB,mBAAA;AAAA,IAAmB;AAAA,EAElD;AACF;"}

package/dist/core-DBxmxwCi.cjs

@@ -1,2 +0,0 @@
-"use strict";async function r(t){const o=process.env.GRAM_API_URL??"https://app.getgram.ai";try{const e=await fetch(o+"/rpc/chatSessions.create",{method:"POST",body:JSON.stringify({embed_origin:t.options.embedOrigin,user_identifier:t.options.userIdentifier,expires_after:t.options.expiresAfter}),headers:{"Content-Type":"application/json","Gram-Project":t.projectSlug,"Gram-Key":t.options.apiKey??process.env.GRAM_API_KEY??""}}),s=await e.text();return{status:e.status,body:s,headers:{"Content-Type":"application/json"}}}catch(e){const s=e instanceof Error?e.message:"Unknown error";return console.error("Failed to create chat session:",e),{status:500,body:JSON.stringify({error:"Failed to create chat session: "+s}),headers:{"Content-Type":"application/json"}}}}exports.createChatSession=r;
-//# sourceMappingURL=core-DBxmxwCi.cjs.map

package/dist/core-DBxmxwCi.cjs.map

@@ -1 +0,0 @@
-{"version":3,"file":"core-DBxmxwCi.cjs","sources":["../src/server/core.ts"],"sourcesContent":null,"names":["createChatSession","request","base","response","body","error","errorMessage"],"mappings":"aA2CA,eAAsBA,EACpBC,EACgC,CAChC,MAAMC,EAAO,QAAQ,IAAI,cAAgB,yBAEzC,GAAI,CACF,MAAMC,EAAW,MAAM,MAAMD,EAAO,2BAA4B,CAC9D,OAAQ,OACR,KAAM,KAAK,UAAU,CACnB,aAAcD,EAAQ,QAAQ,YAC9B,gBAAiBA,EAAQ,QAAQ,eACjC,cAAeA,EAAQ,QAAQ,YAAA,CAChC,EACD,QAAS,CACP,eAAgB,mBAChB,eAAgBA,EAAQ,YACxB,WAAYA,EAAQ,QAAQ,QAAU,QAAQ,IAAI,cAAgB,EAAA,CACpE,CACD,EAEKG,EAAO,MAAMD,EAAS,KAAA,EAE5B,MAAO,CACL,OAAQA,EAAS,OACjB,KAAAC,EACA,QAAS,CAAE,eAAgB,kBAAA,CAAmB,CAElD,OAASC,EAAO,CACd,MAAMC,EACJD,aAAiB,MAAQA,EAAM,QAAU,gBAC3C,eAAQ,MAAM,iCAAkCA,CAAK,EAE9C,CACL,OAAQ,IACR,KAAM,KAAK,UAAU,CACnB,MAAO,kCAAoCC,CAAA,CAC5C,EACD,QAAS,CAAE,eAAgB,kBAAA,CAAmB,CAElD,CACF"}