diff --git a/agent/stockbroker/nodes/tools.tsx b/agent/stockbroker/nodes/tools.tsx
index 9fbe432..d31a7dc 100644
--- a/agent/stockbroker/nodes/tools.tsx
+++ b/agent/stockbroker/nodes/tools.tsx
@@ -1,4 +1,4 @@
-import { StockbrokerState } from "../types";
+import { StockbrokerState, StockbrokerUpdate } from "../types";
 import { ChatOpenAI } from "@langchain/openai";
 import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
 import type ComponentMap from "../../uis/index";
@@ -32,7 +32,7 @@ const STOCKBROKER_TOOLS = [
 export async function callTools(
   state: StockbrokerState,
   config: LangGraphRunnableConfig,
-): Promise<Partial<StockbrokerState>> {
+): Promise<StockbrokerUpdate> {
   const ui = typedUi<typeof ComponentMap>(config);
 
   const message = await llm.bindTools(STOCKBROKER_TOOLS).invoke([
@@ -65,8 +65,7 @@ export async function callTools(
 
   return {
     messages: [message],
-    // TODO: Fix the ui return type.
-    ui: ui.collect as any[],
+    ui: ui.collect,
     timestamp: Date.now(),
   };
 }
diff --git a/agent/stockbroker/types.ts b/agent/stockbroker/types.ts
index ebe3fdb..ab8f43b 100644
--- a/agent/stockbroker/types.ts
+++ b/agent/stockbroker/types.ts
@@ -8,3 +8,4 @@ export const StockbrokerAnnotation = Annotation.Root({
 });
 
 export type StockbrokerState = typeof StockbrokerAnnotation.State;
+export type StockbrokerUpdate = typeof StockbrokerAnnotation.Update;
diff --git a/agent/trip-planner/nodes/classify.ts b/agent/trip-planner/nodes/classify.ts
index 945c056..05e2df8 100644
--- a/agent/trip-planner/nodes/classify.ts
+++ b/agent/trip-planner/nodes/classify.ts
@@ -1,11 +1,11 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { TripPlannerState } from "../types";
+import { TripPlannerState, TripPlannerUpdate } from "../types";
 import { z } from "zod";
 import { formatMessages } from "agent/utils/format-messages";
 
 export async function classify(
   state: TripPlannerState,
-): Promise<Partial<TripPlannerState>> {
+): Promise<TripPlannerUpdate> {
   if (!state.tripDetails) {
     // Can not classify if tripDetails are undefined
     return {};
diff --git a/agent/trip-planner/nodes/extraction.tsx b/agent/trip-planner/nodes/extraction.tsx
index 264e992..3950f52 100644
--- a/agent/trip-planner/nodes/extraction.tsx
+++ b/agent/trip-planner/nodes/extraction.tsx
@@ -1,5 +1,5 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { TripDetails, TripPlannerState } from "../types";
+import { TripDetails, TripPlannerState, TripPlannerUpdate } from "../types";
 import { z } from "zod";
 import { formatMessages } from "agent/utils/format-messages";
 
@@ -43,7 +43,7 @@ function calculateDates(
 
 export async function extraction(
   state: TripPlannerState,
-): Promise<Partial<TripPlannerState>> {
+): Promise<TripPlannerUpdate> {
   const schema = z.object({
     location: z
       .string()
diff --git a/agent/trip-planner/nodes/tools.tsx b/agent/trip-planner/nodes/tools.tsx
index 9d45486..e32e2b4 100644
--- a/agent/trip-planner/nodes/tools.tsx
+++ b/agent/trip-planner/nodes/tools.tsx
@@ -1,4 +1,4 @@
-import { TripPlannerState } from "../types";
+import { TripPlannerState, TripPlannerUpdate } from "../types";
 import { ChatOpenAI } from "@langchain/openai";
 import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
 import type ComponentMap from "../../uis/index";
@@ -48,7 +48,7 @@ const schema = z.object({
 export async function callTools(
   state: TripPlannerState,
   config: LangGraphRunnableConfig,
-): Promise<Partial<TripPlannerState>> {
+): Promise<TripPlannerUpdate> {
   if (!state.tripDetails) {
     throw new Error("No trip details found");
   }
@@ -111,8 +111,7 @@ export async function callTools(
 
   return {
     messages: [response],
-    // TODO: Fix the ui return type.
-    ui: ui.collect as any[],
+    ui: ui.collect,
     timestamp: Date.now(),
   };
 }
diff --git a/agent/trip-planner/types.ts b/agent/trip-planner/types.ts
index 9e546bb..2fc66a4 100644
--- a/agent/trip-planner/types.ts
+++ b/agent/trip-planner/types.ts
@@ -16,3 +16,4 @@ export const TripPlannerAnnotation = Annotation.Root({
 });
 
 export type TripPlannerState = typeof TripPlannerAnnotation.State;
+export type TripPlannerUpdate = typeof TripPlannerAnnotation.Update;
\ No newline at end of file