From c4c3303f1ef57e564857c27cab1e6e1d58a16dd8 Mon Sep 17 00:00:00 2001
From: Joaquim Verges <joaquim.verges@gmail.com>
Date: Tue, 7 Oct 2025 21:38:56 +1300
Subject: [PATCH] [AI-SDK] Add default in-memory session store with custom
 store option

---
 .changeset/eleven-kings-fix.md                |  5 +++++
 .../ai/ai-sdk/components/chat-container.tsx   |  2 +-
 .../ai-sdk-provider/src/exports/thirdweb.ts   |  6 +++++-
 packages/ai-sdk-provider/src/provider.ts      | 21 ++++++++++++-------
 packages/ai-sdk-provider/src/types.ts         |  6 ++++++
 5 files changed, 31 insertions(+), 9 deletions(-)
 create mode 100644 .changeset/eleven-kings-fix.md

diff --git a/.changeset/eleven-kings-fix.md b/.changeset/eleven-kings-fix.md
new file mode 100644
index 00000000000..a9d58657ab1
--- /dev/null
+++ b/.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
diff --git a/apps/playground-web/src/app/ai/ai-sdk/components/chat-container.tsx b/apps/playground-web/src/app/ai/ai-sdk/components/chat-container.tsx
index d5c536cfc47..218fa3128fc 100644
--- a/apps/playground-web/src/app/ai/ai-sdk/components/chat-container.tsx
+++ b/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 />
diff --git a/packages/ai-sdk-provider/src/exports/thirdweb.ts b/packages/ai-sdk-provider/src/exports/thirdweb.ts
index 95ac2c18cc6..e0cb0636f11 100644
--- a/packages/ai-sdk-provider/src/exports/thirdweb.ts
+++ b/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,
diff --git a/packages/ai-sdk-provider/src/provider.ts b/packages/ai-sdk-provider/src/provider.ts
index 96e916ed2f6..8547cbb3801 100644
--- a/packages/ai-sdk-provider/src/provider.ts
+++ b/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);
diff --git a/packages/ai-sdk-provider/src/types.ts b/packages/ai-sdk-provider/src/types.ts
index 85325e6283e..221b6e3ac1f 100644
--- a/packages/ai-sdk-provider/src/types.ts
+++ b/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 {