Merge pull request ag-ui-protocol#241 from ag-ui-protocol/feat/restore-time-travel

Author: ranst91

typescript-sdk/apps/dojo/scripts/generate-content-json.ts

@@ -140,57 +140,60 @@ async function getFeatureFrontendFiles(featureId: string) {
 }
 
 const integrationsFolderPath = '../../../integrations'
-const agentFilesMapper: Record<string, (agentKeys: string[]) => Record<string, string>> = {
+const agentFilesMapper: Record<string, (agentKeys: string[]) => Record<string, string[]>> = {
   'middleware-starter': () => ({
-    agentic_chat: path.join(__dirname, integrationsFolderPath, `/middleware-starter/src/index.ts`)
+    agentic_chat: [path.join(__dirname, integrationsFolderPath, `/middleware-starter/src/index.ts`)]
   }),
   'pydantic-ai': (agentKeys: string[]) => {
     return agentKeys.reduce((acc, agentId) => ({
       ...acc,
-      [agentId]: `https://github.com/pydantic/pydantic-ai/blob/main/examples/pydantic_ai_examples/ag_ui/api/${agentId}.py`
+      [agentId]: [`https://github.com/pydantic/pydantic-ai/blob/main/examples/pydantic_ai_examples/ag_ui/api/${agentId}.py`]
     }), {})
   },
   'server-starter': () => ({
-    agentic_chat: path.join(__dirname, integrationsFolderPath, `/server-starter/server/python/example_server/__init__.py`)
+    agentic_chat: [path.join(__dirname, integrationsFolderPath, `/server-starter/server/python/example_server/__init__.py`)]
   }),
   'server-starter-all-features': (agentKeys: string[]) => {
     return agentKeys.reduce((acc, agentId) => ({
       ...acc,
-      [agentId]: path.join(__dirname, integrationsFolderPath, `/server-starter/server/python/example_server/${agentId}.py`)
+      [agentId]: [path.join(__dirname, integrationsFolderPath, `/server-starter-all-features/server/python/example_server/${agentId}.py`)]
     }), {})
   },
   'mastra': () => ({
-    agentic_chat: path.join(__dirname, integrationsFolderPath, `/mastra/example/src/mastra/agents/weather-agent.ts`)
+    agentic_chat: [path.join(__dirname, integrationsFolderPath, `/mastra/example/src/mastra/agents/weather-agent.ts`)]
   }),
   'mastra-agent-lock': () => ({
-    agentic_chat: path.join(__dirname, integrationsFolderPath, `/mastra/example/src/mastra/agents/weather-agent.ts`)
+    agentic_chat: [path.join(__dirname, integrationsFolderPath, `/mastra/example/src/mastra/agents/weather-agent.ts`)]
   }),
   'vercel-ai-sdk': () => ({
-    agentic_chat: path.join(__dirname, integrationsFolderPath, `/vercel-ai-sdk/src/index.ts`)
+    agentic_chat: [path.join(__dirname, integrationsFolderPath, `/vercel-ai-sdk/src/index.ts`)]
   }),
   'langgraph': (agentKeys: string[]) => {
     return agentKeys.reduce((acc, agentId) => ({
       ...acc,
-      [agentId]: path.join(__dirname, integrationsFolderPath, `/langgraph/examples/agents/${agentId}/agent.py`)
+      [agentId]: [
+        path.join(__dirname, integrationsFolderPath, `/langgraph/examples/python/agents/${agentId}/agent.py`),
+        path.join(__dirname, integrationsFolderPath, `/langgraph/examples/typescript/src/agents/${agentId}/agent.ts`)
+      ]
     }), {})
   },
   'langgraph-fastapi': (agentKeys: string[]) => {
     return agentKeys.reduce((acc, agentId) => ({
       ...acc,
-      [agentId]: path.join(__dirname, integrationsFolderPath, `/langgraph/python/ag_ui_langgraph/examples/agents/${agentId}.py`)
+      [agentId]: [path.join(__dirname, integrationsFolderPath, `/langgraph/examples/python/agents/${agentId}/agent.py`)]
     }), {})
   },
   'agno': () => ({}),
   'llama-index': (agentKeys: string[]) => {
     return agentKeys.reduce((acc, agentId) => ({
       ...acc,
-      [agentId]: path.join(__dirname, integrationsFolderPath, `/llamaindex/server-py/server/routers/${agentId}.py`)
+      [agentId]: [path.join(__dirname, integrationsFolderPath, `/llamaindex/server-py/server/routers/${agentId}.py`)]
     }), {})
   },
   'crewai': (agentKeys: string[]) => {
     return agentKeys.reduce((acc, agentId) => ({
       ...acc,
-      [agentId]: path.join(__dirname, integrationsFolderPath, `/crewai/python/ag_ui_crewai/examples/${agentId}.py`)
+      [agentId]: [path.join(__dirname, integrationsFolderPath, `/crewai/python/ag_ui_crewai/examples/${agentId}.py`)]
     }), {})
   }
 }
@@ -202,14 +205,16 @@ async function runGenerateContent() {
     const agentsPerFeatures = agentConfig.agentKeys
 
     const agentFilePaths = agentFilesMapper[agentConfig.id](agentConfig.agentKeys)
+
     // Per feature, assign all the frontend files like page.tsx as well as all agent files
     for (const featureId of agentsPerFeatures) {
+      const agentFilePathsForFeature = agentFilePaths[featureId] ?? []
       // @ts-expect-error -- redundant error about indexing of a new object.
       result[`${agentConfig.id}::${featureId}`] = [
         // Get all frontend files for the feature
         ...(await getFeatureFrontendFiles(featureId)),
         // Get the agent (python/TS) file
-        await getFile(agentFilePaths[featureId])
+        ...(await Promise.all(agentFilePathsForFeature.map(async f => await getFile(f))))
       ]
     }
   }

typescript-sdk/apps/dojo/src/app/[integrationId]/feature/tool_based_generative_ui/page.tsx

@@ -51,7 +51,7 @@ export default function ToolBasedGenerativeUI({ params }: ToolBasedGenerativeUIP
       agent="tool_based_generative_ui"
     >
       <div
-        className={`${isMobile ? 'h-screen' : 'min-h-full'} w-full relative overflow-hidden`}
+        className={`${isMobile ? 'h-screen' : 'min-h-full flex'} w-full relative overflow-hidden`}
         style={
           {
             // "--copilot-kit-primary-color": "#222",
@@ -336,7 +336,7 @@ function Haiku() {
     <div className="flex h-full w-full">
       {/* Thumbnail List */}
       {Boolean(generatedHaikus.length) && !isMobile && (
-        <div className="w-40 p-4 border-r border-gray-200 overflow-y-auto overflow-x-hidden max-w-1">
+        <div className="w-40 p-4 border-r border-gray-200 overflow-y-auto overflow-x-hidden">
           {generatedHaikus.map((haiku, index) => (
             <div
               key={index}

typescript-sdk/apps/dojo/src/files.json

@@ -2,7 +2,8 @@
 A simple agentic chat flow using LangGraph instead of CrewAI.
 """
 
-from typing import Dict, List, Any, Optional
+from typing import List, Any
+import os
 
 # Updated imports for LangGraph
 from langchain_core.runnables import RunnableConfig
@@ -76,5 +77,16 @@ async def chat_node(state: AgentState, config: RunnableConfig):
 workflow.add_edge(START, "chat_node")
 workflow.add_edge("chat_node", END)
 
+# Conditionally use a checkpointer based on the environment
+# Check for multiple indicators that we're running in LangGraph dev/API mode
+is_fast_api = os.environ.get("LANGGRAPH_FAST_API", "false").lower() == "true"
+
 # Compile the graph
-agentic_chat_graph = workflow.compile()
+if is_fast_api:
+    # For CopilotKit and other contexts, use MemorySaver
+    from langgraph.checkpoint.memory import MemorySaver
+    memory = MemorySaver()
+    graph = workflow.compile(checkpointer=memory)
+else:
+    # When running in LangGraph API/dev, don't use a custom checkpointer
+    graph = workflow.compile()

typescript-sdk/integrations/langgraph/examples/.env.example renamed to typescript-sdk/integrations/langgraph/examples/python/.env.example

@@ -2,9 +2,10 @@
 An example demonstrating agentic generative UI using LangGraph.
 """
 
-import json
 import asyncio
-from typing import Dict, List, Any, Optional, Literal
+from typing import List, Any
+import os
+
 # LangGraph imports
 from langchain_core.runnables import RunnableConfig
 from langgraph.graph import StateGraph, END, START
@@ -189,5 +190,16 @@ async def chat_node(state: AgentState, config: RunnableConfig):
 workflow.add_edge("start_flow", "chat_node")
 workflow.add_edge("chat_node", END)
 
+# Conditionally use a checkpointer based on the environment
+# Check for multiple indicators that we're running in LangGraph dev/API mode
+is_fast_api = os.environ.get("LANGGRAPH_FAST_API", "false").lower() == "true"
+
 # Compile the graph
-graph = workflow.compile()
+if is_fast_api:
+    # For CopilotKit and other contexts, use MemorySaver
+    from langgraph.checkpoint.memory import MemorySaver
+    memory = MemorySaver()
+    graph = workflow.compile(checkpointer=memory)
+else:
+    # When running in LangGraph API/dev, don't use a custom checkpointer
+    graph = workflow.compile()