fix: Refactor graph state to use zod (#10)

* fix: Refactor graph state to use zod

* handle err

src/index.ts

@@ -81,5 +81,6 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   .addEdge("take-action", "progress-plan-step")
   .addEdge("progress-plan-step", "take-action");
 
-export const graph = workflow.compile();
+// Zod types are messed up
+export const graph = workflow.compile() as any;
 graph.name = "LangGraph ReAct MCP";

src/nodes/generate-message.ts

@@ -71,6 +71,6 @@ export async function generateAction(
   ]);
 
   return {
-    messages: response,
+    messages: [response],
   };
 }

src/nodes/generate-plan.ts

@@ -41,7 +41,7 @@ export async function generatePlan(
 
   // No tool calls generated, instead we should just return the messages.
   return {
-    messages: response,
+    messages: [response],
     proposedPlan: [],
     plan: [],
   };

src/types.ts

@@ -1,10 +1,11 @@
 import "@langchain/langgraph/zod";
 import { z } from "zod";
 import {
-  Annotation,
+  addMessages,
   LangGraphRunnableConfig,
-  MessagesAnnotation,
+  Messages,
 } from "@langchain/langgraph";
+import { BaseMessage } from "@langchain/core/messages";
 
 export type PlanItem = {
   /**
@@ -28,24 +29,28 @@ export type TargetRepository = {
   branch?: string;
 };
 
-export const GraphAnnotation = Annotation.Root({
+export const GraphAnnotation = z.object({
-  messages: MessagesAnnotation.spec.messages,
+  messages: z
-  proposedPlan: Annotation<string[]>({
+    .custom<BaseMessage[]>()
-    reducer: (_state, update) => update,
+    .default(() => [])
-    default: () => [],
+    .langgraph.reducer<Messages>((state, update) => addMessages(state, update)),
-  }),
+  proposedPlan: z
-  plan: Annotation<PlanItem[]>({
+    .array(z.string())
-    reducer: (_state, update) => update,
+    .default(() => [])
-    default: () => [],
+    .langgraph.reducer((_state, update) => update),
-  }),
+  plan: z
-  planChangeRequest: Annotation<string | undefined>({
+    .custom<PlanItem[]>()
-    reducer: (_state, update) => update,
+    .default(() => [])
-    default: () => undefined,
+    .langgraph.reducer<PlanItem[]>((_state, update) => update),
-  }),
+  planChangeRequest: z
+    .string()
+    .nullable()
+    .default(() => null)
+    .langgraph.reducer((_state, update) => update),
 });
 
-export type GraphState = typeof GraphAnnotation.State;
+export type GraphState = z.infer<typeof GraphAnnotation>;
-export type GraphUpdate = typeof GraphAnnotation.Update;
+export type GraphUpdate = Partial<GraphState>;
 
 const MODEL_OPTIONS = [
   {