diff --git a/src/index.ts b/src/index.ts
index 493fdfe7..15db59cc 100644
--- a/src/index.ts
+++ b/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";
diff --git a/src/nodes/generate-message.ts b/src/nodes/generate-message.ts
index 71d816c3..54b3a1e9 100644
--- a/src/nodes/generate-message.ts
+++ b/src/nodes/generate-message.ts
@@ -71,6 +71,6 @@ export async function generateAction(
   ]);
 
   return {
-    messages: response,
+    messages: [response],
   };
 }
diff --git a/src/nodes/generate-plan.ts b/src/nodes/generate-plan.ts
index 0a1c7a7a..bad51422 100644
--- a/src/nodes/generate-plan.ts
+++ b/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: [],
   };
diff --git a/src/types.ts b/src/types.ts
index 1b73203c..136a4008 100644
--- a/src/types.ts
+++ b/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({
-  messages: MessagesAnnotation.spec.messages,
-  proposedPlan: Annotation<string[]>({
-    reducer: (_state, update) => update,
-    default: () => [],
-  }),
-  plan: Annotation<PlanItem[]>({
-    reducer: (_state, update) => update,
-    default: () => [],
-  }),
-  planChangeRequest: Annotation<string | undefined>({
-    reducer: (_state, update) => update,
-    default: () => undefined,
-  }),
+export const GraphAnnotation = z.object({
+  messages: z
+    .custom<BaseMessage[]>()
+    .default(() => [])
+    .langgraph.reducer<Messages>((state, update) => addMessages(state, update)),
+  proposedPlan: z
+    .array(z.string())
+    .default(() => [])
+    .langgraph.reducer((_state, update) => update),
+  plan: z
+    .custom<PlanItem[]>()
+    .default(() => [])
+    .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 GraphUpdate = typeof GraphAnnotation.Update;
+export type GraphState = z.infer<typeof GraphAnnotation>;
+export type GraphUpdate = Partial<GraphState>;
 
 const MODEL_OPTIONS = [
   {