fix(cli): correctly handle tool invocation cancellation (#844)

4 files changed, 62 insertions, 0 deletions

diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 180d74bb..cbbbd113 100644
--- a/packages/core/src/core/client.test.ts
+++ b/packages/core/src/core/client.test.ts
@@ -219,4 +219,22 @@ describe('Gemini Client (client.ts)', () => {
       );
     });
   });
+
+  describe('addHistory', () => {
+    it('should call chat.addHistory with the provided content', async () => {
+      const mockChat = {
+        addHistory: vi.fn(),
+      };
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      client['chat'] = Promise.resolve(mockChat as any);
+
+      const newContent = {
+        role: 'user',
+        parts: [{ text: 'New history item' }],
+      };
+      await client.addHistory(newContent);
+
+      expect(mockChat.addHistory).toHaveBeenCalledWith(newContent);
+    });
+  });
 });
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index 0f2a1b8a..8b921ab1 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -58,6 +58,11 @@ export class GeminiClient {
     this.chat = this.startChat();
   }
 
+  async addHistory(content: Content) {
+    const chat = await this.chat;
+    chat.addHistory(content);
+  }
+
   getChat(): Promise<GeminiChat> {
     return this.chat;
   }
diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts
index 6d18ebd9..dbed31b1 100644
--- a/packages/core/src/core/geminiChat.test.ts
+++ b/packages/core/src/core/geminiChat.test.ts
@@ -352,4 +352,34 @@ describe('GeminiChat', () => {
       expect(history[1].parts).toEqual([{ text: 'Visible text' }]);
     });
   });
+
+  describe('addHistory', () => {
+    it('should add a new content item to the history', () => {
+      const newContent: Content = {
+        role: 'user',
+        parts: [{ text: 'A new message' }],
+      };
+      chat.addHistory(newContent);
+      const history = chat.getHistory();
+      expect(history.length).toBe(1);
+      expect(history[0]).toEqual(newContent);
+    });
+
+    it('should add multiple items correctly', () => {
+      const content1: Content = {
+        role: 'user',
+        parts: [{ text: 'Message 1' }],
+      };
+      const content2: Content = {
+        role: 'model',
+        parts: [{ text: 'Message 2' }],
+      };
+      chat.addHistory(content1);
+      chat.addHistory(content2);
+      const history = chat.getHistory();
+      expect(history.length).toBe(2);
+      expect(history[0]).toEqual(content1);
+      expect(history[1]).toEqual(content2);
+    });
+  });
 });
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index 54f74102..47f3f3a6 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -287,6 +287,15 @@ export class GeminiChat {
     return structuredClone(history);
   }
 
+  /**
+   * Adds a new entry to the chat history.
+   *
+   * @param content - The content to add to the history.
+   */
+  addHistory(content: Content): void {
+    this.history.push(content);
+  }
+
   private async *processStreamResponse(
     streamResponse: AsyncGenerator<GenerateContentResponse>,
     inputContent: Content,