From 7185cb38c5c143d390eb22c4eb8035cbf6cf5662 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 +++++---
 js/genkit/tests/chat_test.ts | 65 ++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/js/ai/src/session.ts b/js/ai/src/session.ts
index 2e7f0cf6d..fcdde1392 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);
   }
 
diff --git a/js/genkit/tests/chat_test.ts b/js/genkit/tests/chat_test.ts
index 1d6c148de..e17fddf49 100644
--- a/js/genkit/tests/chat_test.ts
+++ b/js/genkit/tests/chat_test.ts
@@ -660,3 +660,68 @@ describe('preamble', () => {
     ]);
   });
 });
+
+describe('session state', () => {
+  let ai: Genkit;
+
+  beforeEach(() => {
+    ai = genkit({
+      model: 'echoModel',
+    });
+    defineEchoModel(ai);
+  });
+
+  it('properly patches partial state updates', async () => {
+    const session = ai.createSession({
+      initialState: { userName: 'John', color: 'Blue', age: 25 },
+    });
+
+    // Update single property
+    await session.updateState({ color: 'Green' });
+    assert.deepStrictEqual(
+      session.state,
+      { userName: 'John', color: 'Green', age: 25 },
+      'should preserve existing properties when updating a single field'
+    );
+
+    // Update multiple properties
+    await session.updateState({ userName: 'Jane', color: 'red', age: 26 });
+    assert.deepStrictEqual(
+      session.state,
+      { userName: 'Jane', color: 'Green', age: 26 },
+      'should preserve non-updated properties when updating multiple fields'
+    );
+
+    // Update with undefined state
+    const emptySession = ai.createSession();
+    await emptySession.updateState({ newProp: 'value' });
+    assert.deepStrictEqual(
+      emptySession.state,
+      { newProp: 'value' },
+      'should handle updates when initial state is undefined'
+    );
+  });
+
+  it('maintains state across chat messages', async () => {
+    const session = ai.createSession({
+      initialState: { userName: 'Pavel' },
+    });
+
+    const chat = session.chat();
+    await chat.send('hi');
+
+    await session.updateState({ userName: 'Jacob' });
+    assert.deepStrictEqual(
+      session.state,
+      { userName: 'John', status: 'active' },
+      'should preserve state between chat messages'
+    );
+
+    await chat.send('bye');
+    assert.deepStrictEqual(
+      session.state,
+      { userName: 'John', status: 'active' },
+      'state should remain stable after chat messages'
+    );
+  });
+});