[AI-SDK] Add default in-memory session store with custom store option

Author: joaquim-verges

.changeset/eleven-kings-fix.md

@@ -0,0 +1,5 @@
+---
+"@thirdweb-dev/ai-sdk-provider": patch
+---
+
+Default in-memory session store across thirdweb AI instances, option to pass your own

apps/playground-web/src/app/ai/ai-sdk/components/chat-container.tsx

@@ -76,7 +76,7 @@ export function ChatContainer() {
                         return (
                           <Reasoning
                             key={`${message.id}-reasoning-${i}`}
-                            className="w-full"
+                            className="w-full max-w-md"
                             isStreaming={status === "streaming"}
                           >
                             <ReasoningTrigger />

packages/ai-sdk-provider/src/exports/thirdweb.ts

@@ -1,4 +1,8 @@
-export { createThirdwebAI, type ThirdwebProvider } from "../provider.js";
+export {
+  createThirdwebAI,
+  type SessionStore,
+  type ThirdwebProvider,
+} from "../provider.js";
 export type {
   MonitorTransactionInput,
   MonitorTransactionOutput,

packages/ai-sdk-provider/src/provider.ts

@@ -544,7 +544,7 @@ export class ThirdwebProvider implements ProviderV2 {
 
   constructor(config: ThirdwebConfig = {}) {
     this.config = config;
-    this.session = new SessionStore();
+    this.session = config.sessionStore || memorySessionStore;
   }
 
   chat(id?: string, settings: ThirdwebSettings = {}) {
@@ -581,22 +581,29 @@ export class ThirdwebProvider implements ProviderV2 {
   }
 }
 
-class SessionStore {
+export interface SessionStore {
+  getSessionId(chatId: string): string | null;
+  setSessionId(chatId: string, sessionId: string): void;
+  clearSessionId(chatId: string): void;
+}
+
+class InMemorySessionStore implements SessionStore {
   private sessionId: Map<string, string> = new Map();
 
-  getSessionId(chatId: string) {
+  getSessionId(chatId: string): string | null {
     return this.sessionId.get(chatId) || null;
   }
-
-  setSessionId(chatId: string, sessionId: string) {
+  setSessionId(chatId: string, sessionId: string): void {
     this.sessionId.set(chatId, sessionId);
   }
-
-  clearSessionId(chatId: string) {
+  clearSessionId(chatId: string): void {
     this.sessionId.delete(chatId);
   }
 }
 
+// singleton session store used as default if no session store is provided
+const memorySessionStore = new InMemorySessionStore();
+
 // Factory function for easier usage
 export function createThirdwebAI(config: ThirdwebConfig = {}) {
   return new ThirdwebProvider(config);

packages/ai-sdk-provider/src/types.ts

@@ -1,4 +1,5 @@
 import type { UIDataTypes, UIMessage } from "ai";
+import type { SessionStore } from "./provider.js";
 import type {
   MonitorTransactionInput,
   MonitorTransactionOutput,
@@ -26,6 +27,11 @@ export interface ThirdwebConfig {
    * Optional: The user wallet auth token (JWT) for executing transactions
    */
   walletAuthToken?: string;
+  /**
+   * Optional: The session store to use for storing and retrieving thirdweb AI session IDs from ai SDK chat ids.
+   * If not provided, a default in-memory session store will be used.
+   */
+  sessionStore?: SessionStore;
 }
 
 export interface ThirdwebSettings {