Skip to content

Commit 4aab8e0

Browse files
EamonNerbonneEamonNerbonne
authored
Merge pull request #3588 from Kilo-Org/eamon/morecleanups
cleanups in autocomplete loading
2 parents 5e9dad0 + fcf9b13 commit 4aab8e0

File tree

7 files changed

+207
-279
lines changed

7 files changed

+207
-279
lines changed

‎packages/types/src/__tests__/kilocode.test.ts‎

Lines changed: 1 addition & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
11
// npx vitest run src/__tests__/kilocode.test.ts
22

3-
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
3+
import { describe, it, expect, vi, afterEach } from "vitest"
44
import {
55
ghostServiceSettingsSchema,
6-
checkKilocodeBalance,
76
getAppUrl,
87
getKiloUrlFromToken,
98
getExtensionConfigUrl,
@@ -36,118 +35,6 @@ describe("ghostServiceSettingsSchema", () => {
3635
})
3736
})
3837

39-
describe("checkKilocodeBalance", () => {
40-
const mockToken = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJlbnYiOiJwcm9kdWN0aW9uIn0.test"
41-
const mockOrgId = "org-123"
42-
43-
beforeEach(() => {
44-
global.fetch = vi.fn()
45-
})
46-
47-
afterEach(() => {
48-
vi.restoreAllMocks()
49-
})
50-
51-
it("should return true when balance is positive", async () => {
52-
vi.mocked(global.fetch).mockResolvedValueOnce({
53-
ok: true,
54-
json: async () => ({ balance: 100 }),
55-
} as Response)
56-
57-
const result = await checkKilocodeBalance(mockToken)
58-
expect(result).toBe(true)
59-
expect(global.fetch).toHaveBeenCalledWith(
60-
"https://api.kilocode.ai/api/profile/balance",
61-
expect.objectContaining({
62-
headers: expect.objectContaining({
63-
Authorization: `Bearer ${mockToken}`,
64-
}),
65-
}),
66-
)
67-
})
68-
69-
it("should return false when balance is zero", async () => {
70-
vi.mocked(global.fetch).mockResolvedValueOnce({
71-
ok: true,
72-
json: async () => ({ balance: 0 }),
73-
} as Response)
74-
75-
const result = await checkKilocodeBalance(mockToken)
76-
expect(result).toBe(false)
77-
})
78-
79-
it("should return false when balance is negative", async () => {
80-
vi.mocked(global.fetch).mockResolvedValueOnce({
81-
ok: true,
82-
json: async () => ({ balance: -10 }),
83-
} as Response)
84-
85-
const result = await checkKilocodeBalance(mockToken)
86-
expect(result).toBe(false)
87-
})
88-
89-
it("should include organization ID in headers when provided", async () => {
90-
vi.mocked(global.fetch).mockResolvedValueOnce({
91-
ok: true,
92-
json: async () => ({ balance: 100 }),
93-
} as Response)
94-
95-
const result = await checkKilocodeBalance(mockToken, mockOrgId)
96-
expect(result).toBe(true)
97-
expect(global.fetch).toHaveBeenCalledWith(
98-
"https://api.kilocode.ai/api/profile/balance",
99-
expect.objectContaining({
100-
headers: expect.objectContaining({
101-
Authorization: `Bearer ${mockToken}`,
102-
"X-KiloCode-OrganizationId": mockOrgId,
103-
}),
104-
}),
105-
)
106-
})
107-
108-
it("should not include organization ID in headers when not provided", async () => {
109-
vi.mocked(global.fetch).mockResolvedValueOnce({
110-
ok: true,
111-
json: async () => ({ balance: 100 }),
112-
} as Response)
113-
114-
await checkKilocodeBalance(mockToken)
115-
116-
const fetchCall = vi.mocked(global.fetch).mock.calls[0]
117-
expect(fetchCall).toBeDefined()
118-
const headers = (fetchCall![1] as RequestInit)?.headers as Record<string, string>
119-
120-
expect(headers).toHaveProperty("Authorization")
121-
expect(headers).not.toHaveProperty("X-KiloCode-OrganizationId")
122-
})
123-
124-
it("should return false when API request fails", async () => {
125-
vi.mocked(global.fetch).mockResolvedValueOnce({
126-
ok: false,
127-
} as Response)
128-
129-
const result = await checkKilocodeBalance(mockToken)
130-
expect(result).toBe(false)
131-
})
132-
133-
it("should return false when fetch throws an error", async () => {
134-
vi.mocked(global.fetch).mockRejectedValueOnce(new Error("Network error"))
135-
136-
const result = await checkKilocodeBalance(mockToken)
137-
expect(result).toBe(false)
138-
})
139-
140-
it("should handle missing balance field in response", async () => {
141-
vi.mocked(global.fetch).mockResolvedValueOnce({
142-
ok: true,
143-
json: async () => ({}),
144-
} as Response)
145-
146-
const result = await checkKilocodeBalance(mockToken)
147-
expect(result).toBe(false)
148-
})
149-
})
150-
15138
describe("URL functions", () => {
15239
const originalEnv = process.env.KILOCODE_BACKEND_BASE_URL
15340

‎packages/types/src/kilocode/kilocode.ts‎

Lines changed: 0 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import { z } from "zod"
2-
import { ProviderSettings, ProviderSettingsEntry } from "../provider-settings.js"
32

43
declare global {
54
interface Window {
@@ -117,85 +116,3 @@ export function getExtensionConfigUrl(): string {
117116
return "https://api.kilocode.ai/extension-config.json"
118117
}
119118
}
120-
121-
/**
122-
* Check if the Kilocode account has a positive balance
123-
* @param kilocodeToken - The Kilocode JWT token
124-
* @param kilocodeOrganizationId - Optional organization ID to include in headers
125-
* @returns Promise<boolean> - True if balance > 0, false otherwise
126-
*/
127-
export async function checkKilocodeBalance(kilocodeToken: string, kilocodeOrganizationId?: string): Promise<boolean> {
128-
try {
129-
const baseUrl = getKiloBaseUriFromToken(kilocodeToken)
130-
131-
const headers: Record<string, string> = {
132-
Authorization: `Bearer ${kilocodeToken}`,
133-
}
134-
135-
if (kilocodeOrganizationId) {
136-
headers["X-KiloCode-OrganizationId"] = kilocodeOrganizationId
137-
}
138-
139-
const response = await fetch(`${baseUrl}/api/profile/balance`, {
140-
headers,
141-
})
142-
143-
if (!response.ok) {
144-
return false
145-
}
146-
147-
const data = await response.json()
148-
const balance = data.balance ?? 0
149-
return balance > 0
150-
} catch (error) {
151-
console.error("Error checking kilocode balance:", error)
152-
return false
153-
}
154-
}
155-
156-
export const AUTOCOMPLETE_PROVIDER_MODELS = {
157-
mistral: "codestral-latest",
158-
kilocode: "mistralai/codestral-2508",
159-
openrouter: "mistralai/codestral-2508",
160-
bedrock: "mistral.codestral-2508-v1:0",
161-
} as const
162-
export type AutocompleteProviderKey = keyof typeof AUTOCOMPLETE_PROVIDER_MODELS
163-
164-
interface ProviderSettingsManager {
165-
listConfig(): Promise<ProviderSettingsEntry[]>
166-
getProfile(params: { id: string }): Promise<ProviderSettings>
167-
}
168-
169-
export type ProviderUsabilityChecker = (
170-
provider: AutocompleteProviderKey,
171-
providerSettingsManager: ProviderSettingsManager,
172-
) => Promise<boolean>
173-
174-
export const defaultProviderUsabilityChecker: ProviderUsabilityChecker = async (provider, providerSettingsManager) => {
175-
if (provider === "kilocode") {
176-
try {
177-
const profiles = await providerSettingsManager.listConfig()
178-
const kilocodeProfile = profiles.find((p) => p.apiProvider === "kilocode")
179-
180-
if (!kilocodeProfile) {
181-
return false
182-
}
183-
184-
const profile = await providerSettingsManager.getProfile({ id: kilocodeProfile.id })
185-
const kilocodeToken = profile.kilocodeToken
186-
const kilocodeOrgId = profile.kilocodeOrganizationId
187-
188-
if (!kilocodeToken) {
189-
return false
190-
}
191-
192-
return await checkKilocodeBalance(kilocodeToken, kilocodeOrgId)
193-
} catch (error) {
194-
console.error("Error checking kilocode balance:", error)
195-
return false
196-
}
197-
}
198-
199-
// For all other providers, assume they are usable
200-
return true
201-
}

‎src/services/ghost/GhostModel.ts‎

Lines changed: 19 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,9 @@
1-
import {
2-
AUTOCOMPLETE_PROVIDER_MODELS,
3-
defaultProviderUsabilityChecker,
4-
modelIdKeysByProvider,
5-
ProviderSettingsEntry,
6-
} from "@roo-code/types"
1+
import { modelIdKeysByProvider, ProviderSettingsEntry } from "@roo-code/types"
72
import { ApiHandler, buildApiHandler } from "../../api"
83
import { ProviderSettingsManager } from "../../core/config/ProviderSettingsManager"
94
import { OpenRouterHandler } from "../../api/providers"
105
import { ApiStreamChunk } from "../../api/transform/stream"
6+
import { AUTOCOMPLETE_PROVIDER_MODELS, checkKilocodeBalance } from "./utils/kilocode-utils"
117

128
export class GhostModel {
139
private apiHandler: ApiHandler | null = null
@@ -26,50 +22,34 @@ export class GhostModel {
2622

2723
public async reload(providerSettingsManager: ProviderSettingsManager): Promise<boolean> {
2824
const profiles = await providerSettingsManager.listConfig()
29-
const supportedProviders = Object.keys(AUTOCOMPLETE_PROVIDER_MODELS) as Array<
30-
keyof typeof AUTOCOMPLETE_PROVIDER_MODELS
31-
>
3225

3326
this.cleanup()
3427

3528
// Check providers in order, but skip unusable ones (e.g., kilocode with zero balance)
36-
for (const provider of supportedProviders) {
37-
const selectedProfile = profiles.find(
38-
(x): x is typeof x & { apiProvider: string } => x?.apiProvider === provider,
39-
)
40-
if (selectedProfile) {
41-
const isUsable = await defaultProviderUsabilityChecker(provider, providerSettingsManager)
42-
if (!isUsable) continue
43-
44-
this.loadProfile(providerSettingsManager, selectedProfile, provider)
45-
this.loaded = true
46-
return true
29+
for (const [provider, model] of AUTOCOMPLETE_PROVIDER_MODELS) {
30+
const selectedProfile = profiles.find((x) => x?.apiProvider === provider)
31+
if (!selectedProfile) continue
32+
const profile = await providerSettingsManager.getProfile({ id: selectedProfile.id })
33+
34+
if (provider === "kilocode") {
35+
// For all other providers, assume they are usable
36+
if (!profile.kilocodeToken) continue
37+
if (!(await checkKilocodeBalance(profile.kilocodeToken, profile.kilocodeOrganizationId))) continue
4738
}
39+
40+
this.apiHandler = buildApiHandler({ ...profile, [modelIdKeysByProvider[provider]]: model })
41+
42+
if (this.apiHandler instanceof OpenRouterHandler) {
43+
await this.apiHandler.fetchModel()
44+
}
45+
this.loaded = true
46+
return true
4847
}
4948

5049
this.loaded = true // we loaded, and found nothing, but we do not wish to reload
5150
return false
5251
}
5352

54-
public async loadProfile(
55-
providerSettingsManager: ProviderSettingsManager,
56-
selectedProfile: ProviderSettingsEntry,
57-
provider: keyof typeof AUTOCOMPLETE_PROVIDER_MODELS,
58-
): Promise<void> {
59-
const profile = await providerSettingsManager.getProfile({
60-
id: selectedProfile.id,
61-
})
62-
63-
this.apiHandler = buildApiHandler({
64-
...profile,
65-
[modelIdKeysByProvider[provider]]: AUTOCOMPLETE_PROVIDER_MODELS[provider],
66-
})
67-
68-
if (this.apiHandler instanceof OpenRouterHandler) {
69-
await this.apiHandler.fetchModel()
70-
}
71-
}
72-
7353
/**
7454
* Generate response with streaming callback support
7555
*/

‎src/services/ghost/__tests__/GhostModel.spec.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { describe, it, expect, vi, beforeEach } from "vitest"
22
import { GhostModel } from "../GhostModel"
33
import { ProviderSettingsManager } from "../../../core/config/ProviderSettingsManager"
4-
import { AUTOCOMPLETE_PROVIDER_MODELS } from "@roo-code/types"
4+
import { AUTOCOMPLETE_PROVIDER_MODELS } from "../utils/kilocode-utils"
55

66
describe("GhostModel", () => {
77
let mockProviderSettingsManager: ProviderSettingsManager
@@ -15,7 +15,7 @@ describe("GhostModel", () => {
1515

1616
describe("reload", () => {
1717
it("sorts profiles by supportedProviders index order", async () => {
18-
const supportedProviders = Object.keys(AUTOCOMPLETE_PROVIDER_MODELS)
18+
const supportedProviders = [...AUTOCOMPLETE_PROVIDER_MODELS.keys()]
1919
const profiles = [
2020
{ id: "3", name: "profile3", apiProvider: supportedProviders[2] },
2121
{ id: "1", name: "profile1", apiProvider: supportedProviders[0] },
@@ -37,7 +37,7 @@ describe("GhostModel", () => {
3737
})
3838

3939
it("filters out profiles without apiProvider", async () => {
40-
const supportedProviders = Object.keys(AUTOCOMPLETE_PROVIDER_MODELS)
40+
const supportedProviders = [...AUTOCOMPLETE_PROVIDER_MODELS.keys()]
4141
const profiles = [
4242
{ id: "1", name: "profile1", apiProvider: undefined },
4343
{ id: "2", name: "profile2", apiProvider: supportedProviders[0] },
@@ -58,7 +58,7 @@ describe("GhostModel", () => {
5858
})
5959

6060
it("filters out profiles with unsupported apiProvider", async () => {
61-
const supportedProviders = Object.keys(AUTOCOMPLETE_PROVIDER_MODELS)
61+
const supportedProviders = [...AUTOCOMPLETE_PROVIDER_MODELS.keys()]
6262
const profiles = [
6363
{ id: "1", name: "profile1", apiProvider: "unsupported" },
6464
{ id: "2", name: "profile2", apiProvider: supportedProviders[0] },
@@ -90,7 +90,7 @@ describe("GhostModel", () => {
9090
})
9191

9292
it("returns true when profile found", async () => {
93-
const supportedProviders = Object.keys(AUTOCOMPLETE_PROVIDER_MODELS)
93+
const supportedProviders = [...AUTOCOMPLETE_PROVIDER_MODELS.keys()]
9494
const profiles = [{ id: "1", name: "profile1", apiProvider: supportedProviders[0] }] as any
9595

9696
vi.mocked(mockProviderSettingsManager.listConfig).mockResolvedValue(profiles)
@@ -303,7 +303,7 @@ describe("GhostModel", () => {
303303
})
304304

305305
it("returns provider name from API handler when provider is loaded", async () => {
306-
const supportedProviders = Object.keys(AUTOCOMPLETE_PROVIDER_MODELS)
306+
const supportedProviders = [...AUTOCOMPLETE_PROVIDER_MODELS.keys()]
307307
const profiles = [{ id: "1", name: "profile1", apiProvider: supportedProviders[0] }] as any
308308

309309
vi.mocked(mockProviderSettingsManager.listConfig).mockResolvedValue(profiles)

0 commit comments

Comments
 (0)