Skip to content

Commit c85855e

Browse files
committed
Internal codebase indexing
1 parent 2bf8eac commit c85855e

31 files changed

+4113
-29
lines changed

pnpm-lock.yaml

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

src/core/prompts/tools/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,7 @@ export function getToolDescriptionsForMode(
135135
// Conditionally exclude codebase_search if feature is disabled or not configured
136136
if (
137137
!codeIndexManager ||
138+
!codeIndexManager?.isManagedIndexingAvailable || // kilocode_change
138139
!(codeIndexManager.isFeatureEnabled && codeIndexManager.isFeatureConfigured && codeIndexManager.isInitialized)
139140
) {
140141
tools.delete("codebase_search")

src/core/tools/codebaseSearchTool.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -185,8 +185,7 @@ ${jsonResult.results
185185
(result) => `File path: ${result.filePath}
186186
Score: ${result.score}
187187
Lines: ${result.startLine}-${result.endLine}
188-
Code Chunk: ${result.codeChunk}
189-
`,
188+
${result.codeChunk ? `Code Chunk: ${result.codeChunk}\n` : ""}`, // kilocode_change - don't include code chunk managed indexing
190189
)
191190
.join("\n")}`
192191

src/core/webview/ClineProvider.ts

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1358,6 +1358,10 @@ ${prompt}
13581358
}
13591359

13601360
await TelemetryService.instance.updateIdentity(providerSettings.kilocodeToken ?? "") // kilocode_change
1361+
1362+
// kilocode_change start Update code index with new Kilo org props
1363+
await this.updateCodeIndexWithKiloProps()
1364+
// kilocode_change end
13611365
} else {
13621366
await this.updateGlobalState("listApiConfigMeta", await this.providerSettingsManager.listConfig())
13631367
}
@@ -1423,6 +1427,10 @@ ${prompt}
14231427
await this.postStateToWebview()
14241428
await TelemetryService.instance.updateIdentity(providerSettings.kilocodeToken ?? "") // kilocode_change
14251429

1430+
// kilocode_change start Update code index with new Kilo org props
1431+
await this.updateCodeIndexWithKiloProps()
1432+
// kilocode_change end
1433+
14261434
if (providerSettings.apiProvider) {
14271435
this.emit(RooCodeEventName.ProviderProfileChanged, { name, provider: providerSettings.apiProvider })
14281436
}
@@ -3350,4 +3358,80 @@ Here is the project's README to help you get started:\n\n${mcpDetails.readmeCont
33503358
return vscode.Uri.file(filePath).toString()
33513359
}
33523360
}
3361+
3362+
// kilocode_change start
3363+
/**
3364+
* Updates the code index manager with current Kilo org credentials
3365+
* This should be called whenever the API configuration changes
3366+
*/
3367+
private async updateCodeIndexWithKiloProps(): Promise<void> {
3368+
try {
3369+
const { apiConfiguration } = await this.getState()
3370+
3371+
// Only proceed if we have both required credentials
3372+
if (!apiConfiguration.kilocodeToken || !apiConfiguration.kilocodeOrganizationId) {
3373+
return
3374+
}
3375+
3376+
// Get kilocodeTesterWarningsDisabledUntil from context
3377+
const kilocodeTesterWarningsDisabledUntil = this.contextProxy.getValue(
3378+
"kilocodeTesterWarningsDisabledUntil",
3379+
)
3380+
3381+
// Fetch organization settings to check if code indexing is enabled
3382+
const { OrganizationService } = await import("../../services/kilocode/OrganizationService")
3383+
const organization = await OrganizationService.fetchOrganization(
3384+
apiConfiguration.kilocodeToken,
3385+
apiConfiguration.kilocodeOrganizationId,
3386+
kilocodeTesterWarningsDisabledUntil,
3387+
)
3388+
3389+
// Check if code indexing is enabled for this organization
3390+
const codeIndexingEnabled = OrganizationService.isCodeIndexingEnabled(organization)
3391+
3392+
if (!codeIndexingEnabled) {
3393+
this.log("[updateCodeIndexWithKiloProps] Code indexing is disabled for this organization")
3394+
return
3395+
}
3396+
3397+
// Get project ID from Kilo config
3398+
const kiloConfig = await this.getKiloConfig()
3399+
const projectId = kiloConfig?.project?.id
3400+
3401+
if (!projectId) {
3402+
this.log("[updateCodeIndexWithKiloProps] No projectId found in Kilo config, skipping code index update")
3403+
return
3404+
}
3405+
3406+
// Get or create the code index manager for the current workspace
3407+
let codeIndexManager = this.getCurrentWorkspaceCodeIndexManager()
3408+
3409+
// If manager doesn't exist yet, it will be created on first access
3410+
// We need to ensure it's initialized with the context proxy
3411+
if (!codeIndexManager) {
3412+
// Try to get the manager again, which will create it if workspace exists
3413+
const workspacePath = this.cwd
3414+
if (workspacePath) {
3415+
codeIndexManager = CodeIndexManager.getInstance(this.context, workspacePath)
3416+
}
3417+
}
3418+
3419+
if (codeIndexManager) {
3420+
// Set the Kilo org props - code indexing is enabled
3421+
codeIndexManager.setKiloOrgCodeIndexProps({
3422+
kilocodeToken: apiConfiguration.kilocodeToken,
3423+
organizationId: apiConfiguration.kilocodeOrganizationId,
3424+
projectId,
3425+
})
3426+
3427+
// Initialize the manager with context proxy if not already initialized
3428+
if (!codeIndexManager.isInitialized) {
3429+
await codeIndexManager.initialize(this.contextProxy)
3430+
}
3431+
}
3432+
} catch (error) {
3433+
this.log(`Failed to update code index with Kilo props: ${error}`)
3434+
}
3435+
}
3436+
// kilocode_change end
33533437
}
Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import { webviewMessageHandler } from "../webviewMessageHandler"
3+
import type { ClineProvider } from "../ClineProvider"
4+
5+
// Mock the getKiloConfig method
6+
vi.mock("../../../utils/path", () => ({
7+
getWorkspacePath: vi.fn(() => "/test/workspace"),
8+
}))
9+
10+
describe("webviewMessageHandler - requestIndexingStatus with managed indexing", () => {
11+
let mockProvider: any
12+
let mockManager: any
13+
14+
beforeEach(() => {
15+
vi.clearAllMocks()
16+
17+
mockManager = {
18+
getCurrentStatus: vi.fn(() => ({
19+
systemStatus: "Standby",
20+
message: "",
21+
processedItems: 0,
22+
totalItems: 0,
23+
currentItemUnit: "items",
24+
workspacePath: "/test/workspace",
25+
})),
26+
getKiloOrgCodeIndexProps: vi.fn(() => null),
27+
setKiloOrgCodeIndexProps: vi.fn(),
28+
workspacePath: "/test/workspace",
29+
}
30+
31+
mockProvider = {
32+
getCurrentWorkspaceCodeIndexManager: vi.fn(() => mockManager),
33+
getState: vi.fn(async () => ({
34+
apiConfiguration: {
35+
kilocodeToken: "test-token",
36+
kilocodeOrganizationId: "test-org-id",
37+
},
38+
})),
39+
getKiloConfig: vi.fn(async () => ({
40+
project: {
41+
id: "test-project-id",
42+
},
43+
})),
44+
postMessageToWebview: vi.fn(),
45+
log: vi.fn(),
46+
} as unknown as ClineProvider
47+
})
48+
49+
it("should set Kilo org props before getting status when organization credentials are available", async () => {
50+
await webviewMessageHandler(mockProvider, {
51+
type: "requestIndexingStatus",
52+
})
53+
54+
// Verify that setKiloOrgCodeIndexProps was called with correct props
55+
expect(mockManager.setKiloOrgCodeIndexProps).toHaveBeenCalledWith({
56+
kilocodeToken: "test-token",
57+
organizationId: "test-org-id",
58+
projectId: "test-project-id",
59+
})
60+
61+
// Verify that status was retrieved and sent to webview
62+
expect(mockManager.getCurrentStatus).toHaveBeenCalled()
63+
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
64+
type: "indexingStatusUpdate",
65+
values: expect.objectContaining({
66+
systemStatus: "Standby",
67+
}),
68+
})
69+
})
70+
71+
it("should not set Kilo org props if they are already set", async () => {
72+
// Mock that props are already set
73+
mockManager.getKiloOrgCodeIndexProps.mockReturnValue({
74+
kilocodeToken: "test-token",
75+
organizationId: "test-org-id",
76+
projectId: "test-project-id",
77+
})
78+
79+
await webviewMessageHandler(mockProvider, {
80+
type: "requestIndexingStatus",
81+
})
82+
83+
// Verify that setKiloOrgCodeIndexProps was NOT called
84+
expect(mockManager.setKiloOrgCodeIndexProps).not.toHaveBeenCalled()
85+
86+
// Verify that status was still retrieved
87+
expect(mockManager.getCurrentStatus).toHaveBeenCalled()
88+
})
89+
90+
it("should not set Kilo org props if organization credentials are missing", async () => {
91+
mockProvider.getState = vi.fn(async () => ({
92+
apiConfiguration: {
93+
// No kilocodeToken or kilocodeOrganizationId
94+
},
95+
}))
96+
97+
await webviewMessageHandler(mockProvider, {
98+
type: "requestIndexingStatus",
99+
})
100+
101+
// Verify that setKiloOrgCodeIndexProps was NOT called
102+
expect(mockManager.setKiloOrgCodeIndexProps).not.toHaveBeenCalled()
103+
104+
// Verify that status was still retrieved
105+
expect(mockManager.getCurrentStatus).toHaveBeenCalled()
106+
})
107+
108+
it("should not set Kilo org props if project ID is missing", async () => {
109+
mockProvider.getKiloConfig = vi.fn(async () => ({
110+
project: {
111+
// No id
112+
},
113+
}))
114+
115+
await webviewMessageHandler(mockProvider, {
116+
type: "requestIndexingStatus",
117+
})
118+
119+
// Verify that setKiloOrgCodeIndexProps was NOT called
120+
expect(mockManager.setKiloOrgCodeIndexProps).not.toHaveBeenCalled()
121+
122+
// Verify that status was still retrieved
123+
expect(mockManager.getCurrentStatus).toHaveBeenCalled()
124+
})
125+
126+
it("should send error status when no workspace is open", async () => {
127+
mockProvider.getCurrentWorkspaceCodeIndexManager = vi.fn(() => null)
128+
129+
await webviewMessageHandler(mockProvider, {
130+
type: "requestIndexingStatus",
131+
})
132+
133+
// Verify error status was sent
134+
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
135+
type: "indexingStatusUpdate",
136+
values: {
137+
systemStatus: "Error",
138+
message: "orchestrator.indexingRequiresWorkspace",
139+
processedItems: 0,
140+
totalItems: 0,
141+
currentItemUnit: "items",
142+
workerspacePath: undefined,
143+
},
144+
})
145+
})
146+
})

0 commit comments

Comments
 (0)