Browse Source

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.
Wojciech Kordalski 8 months ago
parent
commit
74faacd69d
1 changed files with 33 additions and 2 deletions
  1. 33 2
      src/api/providers/fake-ai.ts

+ 33 - 2
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 {