Merge pull request #310 from ag-ui-protocol/fix/report-state-langgraph

fix: fix picking and reporting of state snapshots

Author: ranst91

typescript-sdk/apps/dojo/src/components/sidebar/sidebar.tsx

@@ -33,7 +33,6 @@ export function Sidebar({ isMobile, onMobileClose }: SidebarProps) {
   const pathname = usePathname();
   const { theme, setTheme } = useTheme();
   const isDarkTheme = theme === "dark"
-  console.log(isDarkTheme)
   const { view, frameworkPickerHidden, viewPickerHidden, featurePickerHidden, setView} = useURLParams();
 
   // Extract the current integration ID from the pathname

typescript-sdk/integrations/langgraph/src/agent.ts

@@ -200,7 +200,7 @@ export class LangGraphAgent extends AbstractAgent {
         checkpointId: fork.checkpoint.checkpoint_id!,
         streamMode,
       }),
-      state: timeTravelCheckpoint,
+      state: timeTravelCheckpoint as ThreadState<State>,
       streamMode,
     };
   }
@@ -486,7 +486,7 @@ export class LangGraphAgent extends AbstractAgent {
       this.endStep()
       this.dispatchEvent({
         type: EventType.STATE_SNAPSHOT,
-        snapshot: this.getStateSnapshot(state.values),
+        snapshot: this.getStateSnapshot(state),
       });
       this.dispatchEvent({
         type: EventType.MESSAGES_SNAPSHOT,
@@ -707,7 +707,7 @@ export class LangGraphAgent extends AbstractAgent {
           this.activeRun!.manuallyEmittedState = event.data;
           this.dispatchEvent({
             type: EventType.STATE_SNAPSHOT,
-            snapshot: this.getStateSnapshot(this.activeRun!.manuallyEmittedState!),
+            snapshot: this.getStateSnapshot({ values: this.activeRun!.manuallyEmittedState! } as ThreadState<State>),
             rawEvent: event,
           });
         }
@@ -767,7 +767,8 @@ export class LangGraphAgent extends AbstractAgent {
     }
   }
 
-  getStateSnapshot(state: State) {
+  getStateSnapshot(threadState: ThreadState<State>) {
+    let state = threadState.values
     const schemaKeys = this.activeRun!.schemaKeys!;
     // Do not emit state keys that are not part of the output schema
     if (schemaKeys?.output) {