[AI SDK Provider] Auto handle session ids

Author: joaquim-verges

.changeset/pink-hands-find.md

@@ -0,0 +1,5 @@
+---
+"@thirdweb-dev/ai-sdk-provider": patch
+---
+
+Auto handle session ids

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

@@ -2,7 +2,10 @@
 
 import { useChat } from "@ai-sdk/react";
 import type { ThirdwebAiMessage } from "@thirdweb-dev/ai-sdk-provider";
-import { DefaultChatTransport } from "ai";
+import {
+  DefaultChatTransport,
+  lastAssistantMessageIsCompleteWithToolCalls,
+} from "ai";
 import { useMemo, useState } from "react";
 import { defineChain, prepareTransaction } from "thirdweb";
 import {
@@ -31,30 +34,19 @@ import { Loader } from "../../../../components/loader";
 import { THIRDWEB_CLIENT } from "../../../../lib/client";
 
 export function ChatContainer() {
-  const [sessionId, setSessionId] = useState("");
-
   const { messages, sendMessage, status, addToolResult } =
     useChat<ThirdwebAiMessage>({
       transport: new DefaultChatTransport({
         api: "/api/chat",
       }),
-      onFinish: ({ message }) => {
-        setSessionId(message.metadata?.session_id ?? "");
-      },
+      sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
     });
   const [input, setInput] = useState("");
 
   const handleSubmit = (e: React.FormEvent<HTMLFormElement>) => {
     e.preventDefault();
     if (input.trim()) {
-      sendMessage(
-        { text: input },
-        {
-          body: {
-            sessionId,
-          },
-        },
-      );
+      sendMessage({ text: input });
       setInput("");
     }
   };
@@ -101,7 +93,6 @@ export function ChatContainer() {
                             addToolResult={addToolResult}
                             sendMessage={sendMessage}
                             toolCallId={part.toolCallId}
-                            sessionId={sessionId}
                           />
                         );
                       case "tool-sign_swap":
@@ -113,7 +104,6 @@ export function ChatContainer() {
                             addToolResult={addToolResult}
                             sendMessage={sendMessage}
                             toolCallId={part.toolCallId}
-                            sessionId={sessionId}
                           />
                         );
                       default:
@@ -159,13 +149,12 @@ type SignTransactionButtonProps = {
       >["input"]
     | undefined;
   addToolResult: ReturnType<typeof useChat<ThirdwebAiMessage>>["addToolResult"];
-  toolCallId: string;
   sendMessage: ReturnType<typeof useChat<ThirdwebAiMessage>>["sendMessage"];
-  sessionId: string;
+  toolCallId: string;
 };
 
 const SignTransactionButton = (props: SignTransactionButtonProps) => {
-  const { input, addToolResult, toolCallId, sendMessage, sessionId } = props;
+  const { input, addToolResult, toolCallId, sendMessage } = props;
   const transactionData: {
     chain_id: number;
     to: string;
@@ -209,21 +198,9 @@ const SignTransactionButton = (props: SignTransactionButtonProps) => {
               chain_id: transaction.chain.id,
             },
           });
-          sendMessage(undefined, {
-            body: {
-              sessionId,
-            },
-          });
         }}
         onError={(error) => {
-          sendMessage(
-            { text: `Transaction failed: ${error.message}` },
-            {
-              body: {
-                sessionId,
-              },
-            },
-          );
+          sendMessage({ text: `Transaction failed: ${error.message}` });
         }}
       >
         Sign Transaction
@@ -242,12 +219,11 @@ type SignSwapButtonProps = {
       >["input"]
     | undefined;
   addToolResult: ReturnType<typeof useChat<ThirdwebAiMessage>>["addToolResult"];
-  toolCallId: string;
   sendMessage: ReturnType<typeof useChat<ThirdwebAiMessage>>["sendMessage"];
-  sessionId: string;
+  toolCallId: string;
 };
 const SignSwapButton = (props: SignSwapButtonProps) => {
-  const { input, addToolResult, toolCallId, sendMessage, sessionId } = props;
+  const { input, addToolResult, toolCallId, sendMessage } = props;
   const transactionData: {
     chain_id: number;
     to: string;
@@ -293,21 +269,9 @@ const SignSwapButton = (props: SignSwapButtonProps) => {
               chain_id: transaction.chain.id,
             },
           });
-          sendMessage(undefined, {
-            body: {
-              sessionId,
-            },
-          });
         }}
         onError={(error) => {
-          sendMessage(
-            { text: `Transaction failed: ${error.message}` },
-            {
-              body: {
-                sessionId,
-              },
-            },
-          );
+          sendMessage({ text: `Transaction failed: ${error.message}` });
         }}
       >
         Sign swap

apps/playground-web/src/app/ai/ai-sdk/page.tsx

@@ -29,7 +29,7 @@ export default function Page() {
         icon={BotIcon}
         title={title}
         description={description}
-        docsLink="https://portal.thirdweb.com/ai/chat?utm_source=playground"
+        docsLink="https://portal.thirdweb.com/ai/chat/ai-sdk?utm_source=playground"
       >
         <ChatExample />
         <div className="h-8" />
@@ -69,12 +69,10 @@ const thirdwebAI = createThirdwebAI({
 });
 
 export async function POST(req: Request) {
-  const { messages, sessionId } = await req.json();
-
+  const { messages, id } = await req.json();
   const result = streamText({
-    model: thirdwebAI.chat({
+    model: thirdwebAI.chat(id, {
       context: {
-        session_id: sessionId, // session id for continuity
         chain_ids: [8453], // optional chain ids
         from: "0x...", // optional from address
         auto_execute_transactions: true, // optional, defaults to false
@@ -86,14 +84,6 @@ export async function POST(req: Request) {
 
   return result.toUIMessageStreamResponse({
     sendReasoning: true, // optional, to send reasoning steps to the client
-    messageMetadata({ part }) {
-      // record session id for continuity
-      if (part.type === "finish-step") {
-        return {
-          session_id: part.response.id,
-        };
-      }
-    },
   });
 }
 
@@ -122,33 +112,19 @@ import { useState } from 'react';
 import { ThirdwebAiMessage } from '@thirdweb-dev/ai-sdk-provider';
 
 export default function Page() {
-  const [sessionId, setSessionId] = useState('');
-  const { messages, sendMessage, status } = useChat<ThirdwebAiMessage>({
+  const { messages, sendMessage } = useChat<ThirdwebAiMessage>({
     transport: new DefaultChatTransport({
       // see server implementation below
       api: '/api/chat',
     }),
-    onFinish: ({ message }) => {
-      // record session id for continuity
-      setSessionId(message.metadata?.session_id ?? '');
-    },
   });
 
-  const send = (message: string) => {
-      sendMessage({ text: message }, {
-        body: {
-          // send session id for continuity
-          sessionId,
-        },
-      });
-  }
-
   return (
     <>
       {messages.map(message => (
         <RenderMessage message={message} />
       ))}
-      <ChatInputBox send={send} />
+      <ChatInputBox send={sendMessage} />
     </>
   );
 }`}

apps/playground-web/src/app/api/chat/route.ts

@@ -11,27 +11,15 @@ const thirdwebAI = createThirdwebAI({
 
 export async function POST(req: Request) {
   const body = await req.json();
-  const { messages, sessionId }: { messages: UIMessage[]; sessionId: string } =
-    body;
+  const { messages, id }: { messages: UIMessage[]; id: string } = body;
 
   const result = streamText({
-    model: thirdwebAI.chat({
-      context: {
-        session_id: sessionId,
-      },
-    }),
+    model: thirdwebAI.chat(id),
     messages: convertToModelMessages(messages),
     tools: thirdwebAI.tools(),
   });
 
   return result.toUIMessageStreamResponse({
     sendReasoning: true,
-    messageMetadata({ part }) {
-      if (part.type === "finish-step") {
-        return {
-          session_id: part.response.id,
-        };
-      }
-    },
   });
 }