From 74faacd69d9d767d0fb6ef9d0c9d2ba01957838b Mon Sep 17 00:00:00 2001
From: Wojciech Kordalski <wkordalski@pm.me>
Date: Wed, 23 Apr 2025 10:42:06 +0200
Subject: [PATCH] FakeAI "controller" object must not be copied (#2463)

The FakeAI object passed by the user must be exactly the same object
that is passed to FakeAIHandler via API configuration.
Unfortunatelly, as the VSCode global state is used as configuration
storage, we lose this property (VSCode global state creates copies of
the object). Also the class of the stored object is lost, so methods
of the object are unavailable.

Therefore, we store the original objects in global variable and
use ID field of FakeAI object to identify the original object.
---
 src/api/providers/fake-ai.ts | 35 +++++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/api/providers/fake-ai.ts b/src/api/providers/fake-ai.ts
index f7509c8b06..68d028338e 100644
--- a/src/api/providers/fake-ai.ts
+++ b/src/api/providers/fake-ai.ts
@@ -4,21 +4,52 @@ import { ApiHandlerOptions, ModelInfo } from "../../shared/api"
 import { ApiStream } from "../transform/stream"
 
 interface FakeAI {
+	/**
+	 * The unique identifier for the FakeAI instance.
+	 * It is used to lookup the original FakeAI object in the fakeAiMap
+	 * when the fakeAI object is read from the VSCode global state.
+	 */
+	readonly id: string
+
+	/**
+	 * A function set by the FakeAIHandler on the FakeAI instance, that removes
+	 * the FakeAI instance from the fakeAIMap when the FakeAI instance is
+	 * no longer needed.
+	 */
+	removeFromCache?: () => void
+
 	createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
 	getModel(): { id: string; info: ModelInfo }
 	countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number>
 	completePrompt(prompt: string): Promise<string>
 }
 
+/**
+ * API providers configuration is stored in the VSCode global state.
+ * Therefore, when a new task is created, the FakeAI object in the configuration
+ * is a new object not related to the original one, but with the same ID.
+ *
+ * We use the ID to lookup the original FakeAI object in the mapping.
+ */
+let fakeAiMap: Map<string, FakeAI> = new Map()
+
 export class FakeAIHandler implements ApiHandler, SingleCompletionHandler {
 	private ai: FakeAI
 
 	constructor(options: ApiHandlerOptions) {
-		if (!options.fakeAi) {
+		const optionsFakeAi = options.fakeAi as FakeAI | undefined
+		if (!optionsFakeAi) {
 			throw new Error("Fake AI is not set")
 		}
 
-		this.ai = options.fakeAi as FakeAI
+		const id = optionsFakeAi.id
+		let cachedFakeAi = fakeAiMap.get(id)
+		if (cachedFakeAi === undefined) {
+			cachedFakeAi = optionsFakeAi
+			cachedFakeAi.removeFromCache = () => fakeAiMap.delete(id)
+			fakeAiMap.set(id, cachedFakeAi)
+		}
+		this.ai = cachedFakeAi
 	}
 
 	async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {