Skip to content

Commit 74faacd

Browse files
authored
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.
1 parent e403a96 commit 74faacd

File tree

1 file changed

+33
-2
lines changed

1 file changed

+33
-2
lines changed

src/api/providers/fake-ai.ts

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,52 @@ import { ApiHandlerOptions, ModelInfo } from "../../shared/api"
44
import { ApiStream } from "../transform/stream"
55

66
interface FakeAI {
7+
/**
8+
* The unique identifier for the FakeAI instance.
9+
* It is used to lookup the original FakeAI object in the fakeAiMap
10+
* when the fakeAI object is read from the VSCode global state.
11+
*/
12+
readonly id: string
13+
14+
/**
15+
* A function set by the FakeAIHandler on the FakeAI instance, that removes
16+
* the FakeAI instance from the fakeAIMap when the FakeAI instance is
17+
* no longer needed.
18+
*/
19+
removeFromCache?: () => void
20+
721
createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
822
getModel(): { id: string; info: ModelInfo }
923
countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number>
1024
completePrompt(prompt: string): Promise<string>
1125
}
1226

27+
/**
28+
* API providers configuration is stored in the VSCode global state.
29+
* Therefore, when a new task is created, the FakeAI object in the configuration
30+
* is a new object not related to the original one, but with the same ID.
31+
*
32+
* We use the ID to lookup the original FakeAI object in the mapping.
33+
*/
34+
let fakeAiMap: Map<string, FakeAI> = new Map()
35+
1336
export class FakeAIHandler implements ApiHandler, SingleCompletionHandler {
1437
private ai: FakeAI
1538

1639
constructor(options: ApiHandlerOptions) {
17-
if (!options.fakeAi) {
40+
const optionsFakeAi = options.fakeAi as FakeAI | undefined
41+
if (!optionsFakeAi) {
1842
throw new Error("Fake AI is not set")
1943
}
2044

21-
this.ai = options.fakeAi as FakeAI
45+
const id = optionsFakeAi.id
46+
let cachedFakeAi = fakeAiMap.get(id)
47+
if (cachedFakeAi === undefined) {
48+
cachedFakeAi = optionsFakeAi
49+
cachedFakeAi.removeFromCache = () => fakeAiMap.delete(id)
50+
fakeAiMap.set(id, cachedFakeAi)
51+
}
52+
this.ai = cachedFakeAi
2253
}
2354

2455
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {

0 commit comments

Comments
 (0)