From 0c52130eb162238522846a69a374fd1fcd750fa9 Mon Sep 17 00:00:00 2001
From: Jacob Cable <jacobcable94@gmail.com>
Date: Mon, 23 Dec 2024 12:44:57 +0000
Subject: [PATCH] fix(js/ai): make updateState patch state instead of replace
 it

---
 js/ai/src/session.ts | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/js/ai/src/session.ts b/js/ai/src/session.ts
index 2e7f0cf6dd..fcdde13922 100644
--- a/js/ai/src/session.ts
+++ b/js/ai/src/session.ts
@@ -85,16 +85,22 @@ export class Session<S = any> {
   }
 
   /**
-   * Update session state data.
+   * Update session state data by patching the existing state.
+   * @param data Partial state update that will be merged with existing state
    */
-  async updateState(data: S): Promise<void> {
+  async updateState(data: Partial<S>): Promise<void> {
     let sessionData = this.sessionData;
     if (!sessionData) {
       sessionData = {} as SessionData<S>;
     }
-    sessionData.state = data;
-    this.sessionData = sessionData;
 
+    // Merge the new data with existing state
+    sessionData.state = {
+      ...sessionData.state,
+      ...data,
+    } as S;
+
+    this.sessionData = sessionData;
     await this.store.save(this.id, sessionData);
   }