fix(apps): add session agent kit api

Author: CocDap

apps/playground/app/chat/page.tsx

@@ -7,7 +7,7 @@ import { Badge } from "@/components/ui/badge"
 import { ScrollArea } from "@/components/ui/scroll-area"
 import { Send, Bot } from "lucide-react"
 import Sidebar from "@/components/sidebar"
-import { useAgentStore } from "@/stores/agent-store"
+import { useAgentStore, useAgentRestore } from "@/stores/agent-store"
 
 
 interface ChatMessage {
@@ -35,6 +35,9 @@ export default function ChatPage() {
 
   const [chatMessages, setChatMessages] = useState<ChatMessage[]>([])
 
+  // Restore agent session on page load
+  useAgentRestore()
+
   // Initialize client-side state
   useEffect(() => {
     setIsClient(true)

apps/playground/app/developer/page.tsx

@@ -16,7 +16,7 @@ import {
   CollapsibleTrigger,
 } from "@/components/ui/collapsible"
 import { ChevronDown } from "lucide-react"
-import { useAgentStore } from "@/stores/agent-store"
+import { useAgentStore, useAgentRestore } from "@/stores/agent-store"
 
 type ToolLike = { name?: string; description?: string; schema?: any; schemaJson?: any; call: (args: any) => Promise<any> }
 type EndpointKey = "assets" | "swap" | "bifrost" | "staking"
@@ -49,6 +49,9 @@ export default function DeveloperPage() {
   const [formData, setFormData] = useState<Record<string, any>>({})
   const [parsedSchema, setParsedSchema] = useState<any>(null)
 
+  // Restore agent session on page load
+  useAgentRestore()
+
   // Initialize tools when agentKit is available
   useEffect(() => {
     const initializeTools = async () => {

apps/playground/stores/agent-store.ts

@@ -1,5 +1,6 @@
 import { create } from 'zustand'
 import { persist } from 'zustand/middleware'
+import React from 'react'
 
 interface AgentConfig {
   privateKey: string
@@ -18,22 +19,50 @@ interface AgentState {
   isInitialized: boolean
   isInitializing: boolean
 
+  // Session management
+  sessionExpiry: number | null
+  
   // Actions
   setConfig: (config: AgentConfig) => void
   initializeAgent: () => Promise<void>
   resetAgent: () => void
   setInitializing: (loading: boolean) => void
+  checkSession: () => boolean
+  restoreAgent: () => Promise<void>
 }
 
 export const useAgentStore = create<AgentState>()(
   persist(
-    (set: any, get: any) => ({
+    (set: any, get: any) => {
+      // Check session on store initialization
+      const checkAndRestoreSession = () => {
+        const { sessionExpiry, isInitialized, config } = get()
+        if (sessionExpiry && isInitialized && config) {
+          const now = Date.now()
+          if (now > sessionExpiry) {
+            // Session expired, reset
+            set({
+              agentKit: null,
+              isInitialized: false,
+              sessionExpiry: null
+            })
+            return false
+          }
+          // Session is valid, but we need to reinitialize the agentKit
+          // since it can't be serialized
+          return true
+        }
+        return false
+      }
+
+      return {
       // Initial state
       config: null,
       isConfigured: false,
       agentKit: null,
       isInitialized: false,
       isInitializing: false,
+      sessionExpiry: null,
 
       // Actions
       setConfig: (config: AgentConfig) => {
@@ -47,6 +76,67 @@ export const useAgentStore = create<AgentState>()(
         set({ isInitializing })
       },
 
+      checkSession: () => {
+        const { sessionExpiry, isInitialized } = get()
+        if (!sessionExpiry || !isInitialized) return false
+        
+        const now = Date.now()
+        if (now > sessionExpiry) {
+          // Session expired, reset agent
+          set({
+            agentKit: null,
+            isInitialized: false,
+            sessionExpiry: null
+          })
+          return false
+        }
+        return true
+      },
+
+      restoreAgent: async () => {
+        const { config, sessionExpiry } = get()
+        if (!config || !sessionExpiry) return
+
+        const now = Date.now()
+        if (now > sessionExpiry) {
+          // Session expired
+          set({
+            agentKit: null,
+            isInitialized: false,
+            sessionExpiry: null
+          })
+          return
+        }
+
+        try {
+          set({ isInitializing: true })
+          
+          const { PolkadotAgentKit } = await import("@polkadot-agent-kit/sdk")
+          
+          const agentKit = new PolkadotAgentKit({
+            privateKey: config.privateKey,
+            keyType: config.keyType,
+            chains: config.chains as any,
+          })
+
+          await agentKit.initializeApi()
+
+          set({ 
+            agentKit,
+            isInitialized: true,
+            isInitializing: false
+          })
+
+          console.log('[AgentStore] Agent restored successfully')
+        } catch (error) {
+          console.error('[AgentStore] Failed to restore agent:', error)
+          set({ 
+            isInitialized: false,
+            isInitializing: false
+          })
+        }
+      },
+
       initializeAgent: async () => {
         const { config } = get()
         if (!config || !config.isConfigured) {
@@ -66,10 +156,14 @@ export const useAgentStore = create<AgentState>()(
 
           await agentKit.initializeApi()
 
+          // Set session expiry to 15 minutes from now
+          const sessionExpiry = Date.now() + (15 * 60 * 1000) // 15 minutes
+          
           set({ 
             agentKit,
             isInitialized: true,
-            isInitializing: false
+            isInitializing: false,
+            sessionExpiry
           })
 
           console.log('[AgentStore] Agent initialized successfully')
@@ -89,16 +183,31 @@ export const useAgentStore = create<AgentState>()(
           isConfigured: false,
           agentKit: null,
           isInitialized: false,
-          isInitializing: false
+          isInitializing: false,
+          sessionExpiry: null
         })
       }
+      }
     }),
     {
       name: 'polkadot-agent-store',
       partialize: (state: any) => ({
         config: state.config,
-        isConfigured: state.isConfigured
+        isConfigured: state.isConfigured,
+        isInitialized: state.isInitialized,
+        sessionExpiry: state.sessionExpiry
       })
     }
   )
-)
+
+// Hook to automatically restore agent on page load
+export const useAgentRestore = () => {
+  const { isInitialized, agentKit, restoreAgent, checkSession } = useAgentStore()
+  
+  React.useEffect(() => {
+    // Check if we have a valid session but no agent instance
+    if (checkSession() && !agentKit) {
+      restoreAgent()
+    }
+  }, [])
+}