Code Indexing Support for WSO2 Integrator Copilot

workspaces/ballerina/ballerina-core/src/interfaces/extended-lang-client.ts

@@ -613,7 +613,7 @@ export interface ExecutorPositions {
     executorPositions?: ExecutorPosition[];
 }
 
-// Test Manager related interfaces 
+// Test Manager related interfaces
 
 export interface TestsDiscoveryRequest {
     projectPath: string;
@@ -2055,6 +2055,14 @@ export interface ProjectArtifacts {
     artifacts: Artifacts;
 }
 
+export interface CodeMapRequest {
+    projectPath: string;
+}
+
+export interface CodeMapResponse {
+    content: string;
+}
+
 export interface ProjectInfoRequest {
     projectPath: string;
 }
@@ -2168,6 +2176,7 @@ export interface ExtendedLangClientInterface extends BIInterface {
     updateStatusBar(): void;
     getDidOpenParams(): DidOpenParams;
     getProjectArtifacts(params: ProjectArtifactsRequest): Promise<ProjectArtifacts>;
+    getCodeMap(params: CodeMapRequest): Promise<CodeMapResponse>;
     getProjectInfo(params: ProjectInfoRequest): Promise<ProjectInfo>;
     getSimpleTypeOfExpression(params: GetSimpleTypeOfExpressionRequest): Promise<GetSimpleTypeOfExpressionResponse>;
     openConfigToml(params: OpenConfigTomlRequest): Promise<void>;

workspaces/ballerina/ballerina-extension/src/core/extended-language-client.ts

@@ -277,6 +277,8 @@ import {
     ClausePositionRequest,
     SemanticDiffRequest,
     SemanticDiffResponse,
+    CodeMapRequest,
+    CodeMapResponse,
     ConvertExpressionRequest,
     ConvertExpressionResponse,
     IntrospectDatabaseRequest,
@@ -494,6 +496,7 @@ enum EXTENDED_APIS {
     COPILOT_ALL_LIBRARIES = 'copilotLibraryManager/getLibrariesList',
     COPILOT_FILTER_LIBRARIES = 'copilotLibraryManager/getFilteredLibraries',
     COPILOT_SEARCH_LIBRARIES = 'copilotLibraryManager/getLibrariesBySearch',
+    COPILOT_GET_CODE_MAP = 'designModelService/codeMap',
     GET_MIGRATION_TOOLS = 'projectService/getMigrationTools',
     TIBCO_TO_BI = 'projectService/importTibco',
     MULE_TO_BI = 'projectService/importMule',
@@ -1509,6 +1512,10 @@ export class ExtendedLangClient extends LanguageClient implements ExtendedLangCl
         return this.sendRequest<SemanticDiffResponse>(EXTENDED_APIS.BI_GET_SEMANTIC_DIFF, params);
     }
 
+    async getCodeMap(params: CodeMapRequest): Promise<CodeMapResponse> {
+        return this.sendRequest<CodeMapResponse>(EXTENDED_APIS.COPILOT_GET_CODE_MAP, params);
+    }
+
     // <------------ BI APIS END --------------->
 
 

workspaces/ballerina/ballerina-extension/src/features/ai/activator.ts

@@ -86,6 +86,21 @@ export function activateAIFeatures(ballerinaExternalInstance: BallerinaExtension
                     workspacePath: params.projectPath
                 };
 
+                // Fetch Code Map for the test project
+                let codeMapMarkdown: string | undefined;
+                try {
+                    const codeMapResponse = await langClient.getCodeMap({ projectPath: params.projectPath });
+                    const markdown = codeMapResponse?.content;
+                    if (markdown) {
+                        codeMapMarkdown = markdown;
+                        console.log(`[Test Mode] Code Map fetched for project ${params.projectPath} (${markdown.length} chars)`);
+                    } else {
+                        console.log(`[Test Mode] Code Map response was empty for project ${params.projectPath}`);
+                    }
+                } catch (err) {
+                    console.warn(`[Test Mode] Failed to fetch Code Map, continuing without it:`, err);
+                }
+
                 // Create config using new AICommandConfig pattern
                 const config: AICommandConfig<GenerateAgentCodeRequest> = {
                     executionContext: ctx,
@@ -95,6 +110,7 @@ export function activateAIFeatures(ballerinaExternalInstance: BallerinaExtension
                     params,
                     // No chat storage in test mode
                     chatStorage: undefined,
+                    codeMapMarkdown,
                     // Immediate cleanup (AI_TEST_ENV prevents actual deletion)
                     lifecycle: {
                         cleanupStrategy: 'immediate'

workspaces/ballerina/ballerina-extension/src/features/ai/agent/AgentExecutor.ts

@@ -266,11 +266,20 @@ export class AgentExecutor extends AICommandExecutor<GenerateAgentCodeRequest> {
                 workingDirectory: workspaceId,
             };
 
+
             // Resolve model and login method
             const loginMethod = await getLoginMethod();
             const model = await getAnthropicClient(ANTHROPIC_SONNET_4);
 
-            const userMessageContent = getUserPrompt(params, tempProjectPath, projects);
+            // Code Map is fetched at query submission time in index.ts generateAgent
+            const codeMapMarkdown = this.config.codeMapMarkdown;
+            if (codeMapMarkdown) {
+                console.log(`[AgentExecutor] Code Map included in LLM prompt (${codeMapMarkdown.length} chars)`);
+            } else {
+                console.log(`[AgentExecutor] No Code Map available — sending prompt without Code Map`);
+            }
+
+            const userMessageContent = getUserPrompt(params, tempProjectPath, projects, codeMapMarkdown);
 
             // Estimate fixed overhead (system prompt + codebase) to decide if compaction is viable
             const systemPromptText = getSystemPrompt(projects, params.operationType);
@@ -282,6 +291,7 @@ export class AgentExecutor extends AICommandExecutor<GenerateAgentCodeRequest> {
                 warnCompactionDisabledOnce(projectRootPath, this.config.eventHandler);
             }
 
+
             // 3. Add generation to chat storage (if enabled)
             this.addGeneration(params.usecase, {
                 isPlanMode: params.isPlanMode,
@@ -296,7 +306,6 @@ export class AgentExecutor extends AICommandExecutor<GenerateAgentCodeRequest> {
             // 5. Build LLM messages with history
             const historyMessages = populateHistoryForAgent(chatHistory);
             const cacheOptions = await getProviderCacheControl();
-
             const allMessages: ModelMessage[] = [
                 {
                     role: "system",

workspaces/ballerina/ballerina-extension/src/features/ai/agent/index.ts

@@ -117,6 +117,31 @@ export async function generateAgent(params: GenerateAgentCodeRequest): Promise<b
             }
         );
 
+// ======================================
+// Code Map Generation for Agent Context
+// ======================================
+
+        // Fetch Code Map at query submission time to capture the current project state.
+        // TODO: The language server does not capture file delete events yet. Once that is fixed,
+        // we can use incremental updates — fetching the Code Map only for changed files and
+        // merging into the existing bal.md, instead of regenerating the full Code Map each time.
+        const langClient = StateMachine.context().langClient;
+        const workspacePath = StateMachine.context().workspacePath || StateMachine.context().projectPath;
+        if (langClient) {
+            try {
+                const codeMapResponse = await langClient.getCodeMap({
+                    projectPath: workspacePath,
+                });
+
+                const codeMapMarkdown = codeMapResponse?.content;
+                if (codeMapMarkdown) {
+                    config.codeMapMarkdown = codeMapMarkdown;
+                }
+            } catch (err) {
+                console.warn('[generateAgent] Failed to fetch Code Map, continuing without it:', err);
+            }
+        }
+
         await new AgentExecutor(config).run();
 
         return true;

workspaces/ballerina/ballerina-extension/src/features/ai/agent/prompts.ts

@@ -22,6 +22,8 @@ import { FILE_BATCH_EDIT_TOOL_NAME, FILE_READ_TOOL_NAME, FILE_SINGLE_EDIT_TOOL_N
 import { CONNECTOR_GENERATOR_TOOL } from "./tools/connector-generator";
 import { CONFIG_COLLECTOR_TOOL } from "./tools/config-collector";
 import { CLARIFY_TOOL } from "./tools/clarify";
+import { GREP_TOOL_NAME } from "./tools/grep";
+import { GLOB_TOOL_NAME } from "./tools/glob";
 import { TEST_RUNNER_TOOL_NAME } from "./tools/test-runner";
 import { getLanglibInstructions } from "../utils/libs/langlibs";
 import { formatCodebaseStructure, formatCodeContext } from "./utils";
@@ -53,7 +55,7 @@ Create a very high-level and concise design plan for the given user requirement.
 You must use ${TASK_WRITE_TOOL_NAME} tool to create and manage tasks.
 This plan will be visible to the user and the execution will be guided on the tasks you create.
 
-- Break down the implementation into specific, actionable tasks. Each task should be concise and high level as they are visible to a very high level user. 
+- Break down the implementation into specific, actionable tasks. Each task should be concise and high level as they are visible to a very high level user.
 - Each task should have a type. This type will be used to guide the user through the generation proccess.
 - Track each task as you work through them and mark tasks as you start and complete them
 
@@ -114,7 +116,7 @@ Plan the implementation approach in your reasoning. Keep output minimal — no d
 Identify the libraries required to implement the user requirement. Use ${LIBRARY_SEARCH_TOOL} to discover relevant libraries, then use ${LIBRARY_GET_TOOL} to fetch their full details.
 
 ### Step 3: Write the code
-Write/modify the Ballerina code to implement the user requirement. Use the ${FILE_BATCH_EDIT_TOOL_NAME}, ${FILE_SINGLE_EDIT_TOOL_NAME}, ${FILE_WRITE_TOOL_NAME} tools to write/modify the code. 
+Write/modify the Ballerina code to implement the user requirement. Use the ${FILE_BATCH_EDIT_TOOL_NAME}, ${FILE_SINGLE_EDIT_TOOL_NAME}, ${FILE_WRITE_TOOL_NAME} tools to write/modify the code.
 
 ### Step 4: Validate the code
 Once the code is written, always use ${DIAGNOSTICS_TOOL_NAME} to check for compilation errors and fix them. You may call it multiple times after making changes.
@@ -143,8 +145,6 @@ When generating Ballerina code strictly follow these syntax and structure guidel
 - A submodule MUST BE imported before being used. The import statement should only contain the package name and submodule name. For package my_pkg, folder structure generated/fooApi, the import should be \`import my_pkg.fooApi;\`.
 - For GraphQL service related queries, if the user hasn't specified their own GraphQL Schema, write the proposed GraphQL schema for the user query right after the explanation before generating the Ballerina code. Use the same names as the GraphQL Schema when defining record types.
 - Some libaries has instructions fields in their API documentation. Follow those instructions strictly when using those libraries.
-- You should only generate tests if the user explicitly asks for them in the query. You must use the 'ballerina/test' and whatever services associated when writing tests. Respect the instructions field in ballerina/test library and testGenerationInstruction field in whatever library associated with the service in the library API documentation when writing tests.
-- For workflow-based requirements involving long-running processes, state management, or orchestration of multiple steps, use the 'ballerina/workflow' module.
 - When writing tests, use the 'ballerina/test' module and any service-specific test libraries. Respect the instructions field in ballerina/test library and the testGenerationInstruction field in the associated service library API documentation when writing tests.
 - Some libraries may contain Readme field. This is generic information about the library. Avoid following links from the readme contents.
 ${getLanglibInstructions()}
@@ -176,14 +176,27 @@ ${getLanglibInstructions()}
 - Mention types EXPLICITLY in variable declarations and foreach statements. (Avoid var at all costs)
 - To narrow down a union type(or optional type), always declare a separate variable and then use that variable in the if condition.
 
-## File modifications
+# Codebase Exploration
+- When the user submits a query, you will receive either **Codebase High Level Overview** or **Complete Structure of the Codebase**. Identify which one you have received before proceeding.
+- If you received Codebase High Level Overview, use it as a navigation map to locate the relevant components to the user query, in the codebase, but the actual source must be read separately when needed.
+- Codebase High Level Overview lists, for each Ballerina file, all of its components (imports, configurables, variables, types, functions, services, listeners, classes) with their signatures and line ranges, but excludes implementation bodies, test files, and resource files.
+- If you receive complete structure of the codebase, it contains the complete source of all .bal files (test and resource files excluded) provided directly in your context.
+
+## Context Retrieval
+- Explore the codebase with ${GREP_TOOL_NAME}, ${FILE_READ_TOOL_NAME}, and ${GLOB_TOOL_NAME}, and keep exploring until you have all the context required to answer confidently.
+
+### Rules for exploration
+- **DO NOT** guess the implementation based on signatures from Codebase High Level Overview or excerpts you retrieved. Always read the actual source code before using any information about a component in the codebase. This is critical to avoid hallucinations and wrong assumptions.
+- When you update or write code, Codebase High Level Overview will become outdated.
+
+## File Modifications and Component Modifications
 - You must apply changes to the existing source code using the provided ${[
             FILE_BATCH_EDIT_TOOL_NAME,
             FILE_SINGLE_EDIT_TOOL_NAME,
             FILE_WRITE_TOOL_NAME,
         ].join(
             ", "
-        )} tools. The complete existing source code will be provided in the <existing_code> section of the user prompt.
+        )} tools. The complete existing source code will be provided in the <codebase_structure> section of the user prompt.
 - When making replacements inside an existing file, provide the **exact old string** and the **exact new string** with all newlines, spaces, and indentation, being mindful to replace nearby occurrences together to minimize the number of tool calls.
 - Do NOT create a new markdown file to document each change or summarize your work unless specifically requested by the user.
 - Do not manually add/modify Dependencies.toml. For Config.toml configuration management, use ${CONFIG_COLLECTOR_TOOL}.
@@ -241,13 +254,21 @@ System context:
 </system-reminder>`;
 }
 
-export function getUserPrompt(params: GenerateAgentCodeRequest, tempProjectPath: string, projects: ProjectSource[]) {
+export function getUserPrompt(params: GenerateAgentCodeRequest, tempProjectPath: string, projects: ProjectSource[], codeMapMarkdown?: string) {
     const content = [];
 
-    content.push({
-        type: 'text' as const,
-        text: formatCodebaseStructure(projects, tempProjectPath)
-    });
+    // Add codebase high level summary if available, otherwise fall back to full project structure
+    if (codeMapMarkdown) {
+        content.push({
+            type: 'text' as const,
+            text: `<Codebase High Level Summary>\n${codeMapMarkdown}\n</Codebase High Level Summary>`
+        });
+    } else {
+        content.push({
+            type: 'text' as const,
+            text: formatCodebaseStructure(projects, tempProjectPath)
+        });
+    }
 
     // Add code context if available
     if (params.codeContext) {
@@ -262,25 +283,16 @@ export function getUserPrompt(params: GenerateAgentCodeRequest, tempProjectPath:
 
     // Add file attachments if available
     if (params.fileAttachmentContents && params.fileAttachmentContents.length > 0) {
-        for (const attachment of params.fileAttachmentContents) {
-            if (attachment.mimeType?.startsWith('image/')) {
-                // Add image attachment
-                content.push({
-                    type: 'image' as const,
-                    image: attachment.content,
-                });
-            } else {
-                // Add text file attachment
-                const attachmentsText = params.fileAttachmentContents.map((attachment) =>
-                    `## File: ${attachment.fileName}\n\`\`\`\n${attachment.content}\n\`\`\``).join('\n\n');
-                content.push({
-                    type: 'text' as const,
-                    text: `<User Attachments>
-                            ${attachmentsText}
-                        </User Attachments>`
-                });
-            }
-        }
+        const attachmentsText = params.fileAttachmentContents.map((attachment) =>
+            `## File: ${attachment.fileName}\n\`\`\`\n${attachment.content}\n\`\`\``
+        ).join('\n\n');
+
+        content.push({
+            type: 'text' as const,
+            text: `<User Attachments>
+${attachmentsText}
+</User Attachments>`
+        });
     }
 
     const queryParts = [params.usecase];
@@ -336,5 +348,4 @@ function getNPSuffix(projects: ProjectSource[], op?: OperationType): string {
         basePrompt += getRequirementAnalysisTestGenPrefix(extractResourceDocumentContent(flattenProjectToFiles(projects)));
     }
     return basePrompt;
-}
-
+}

workspaces/ballerina/ballerina-extension/src/features/ai/agent/tool-registry.ts

@@ -55,6 +55,8 @@ import { RunningServicesManager } from './tools/running-service-manager';
 import { createHurlTool, HURL_TOOL_NAME } from './tools/hurl-tool';
 import { createWebSearchTool, WEB_SEARCH_TOOL_NAME, createWebFetchTool, WEB_FETCH_TOOL_NAME } from './tools/web-tools';
 import { createClarifyTool, CLARIFY_TOOL } from './tools/clarify';
+import { createGrepTool, createGrepExecute, GREP_TOOL_NAME } from './tools/grep';
+import { createGlobTool, createGlobExecute, GLOB_TOOL_NAME } from './tools/glob';
 
 export interface ToolRegistryOptions {
     eventHandler: CopilotEventHandler;
@@ -123,6 +125,12 @@ export function createToolRegistry(opts: ToolRegistryOptions) {
         [FILE_READ_TOOL_NAME]: createReadTool(
             createReadExecute(eventHandler, tempProjectPath)
         ),
+        [GREP_TOOL_NAME]: createGrepTool(
+            createGrepExecute(eventHandler, tempProjectPath)
+        ),
+        [GLOB_TOOL_NAME]: createGlobTool(
+            createGlobExecute(eventHandler, tempProjectPath)
+        ),
         [DIAGNOSTICS_TOOL_NAME]: createDiagnosticsTool(tempProjectPath, eventHandler),
         [TEST_RUNNER_TOOL_NAME]: createTestRunnerTool(tempProjectPath, eventHandler, modifiedFiles, allModifiedFiles, ctx),
         // Migration source tools — registered only when a source project path is available