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.
This commit is contained in:
Wojciech Kordalski 2025-04-23 10:42:06 +02:00 committed by GitHub
parent e403a96e66
commit 74faacd69d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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 {