diff --git a/src/api/providers/fake-ai.ts b/src/api/providers/fake-ai.ts index f7509c8b066..68d028338e7 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): Promise completePrompt(prompt: string): Promise } +/** + * 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 = 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 {