wso2 · yasithrashan · May 30, 2026 · May 30, 2026 · May 30, 2026 · coderabbitai
@@ -613,7 +613,7 @@ export interface ExecutorPositions {
     executorPositions?: ExecutorPosition[];
 }
 
-// Test Manager related interfaces 
+// Test Manager related interfaces
 
 export interface TestsDiscoveryRequest {
     projectPath: string;
@@ -2030,6 +2030,14 @@ export interface ProjectArtifacts {
     artifacts: Artifacts;
 }
 
+export interface CodeMapRequest {
+    projectPath: string;
+}
+
+export interface CodeMapResponse {
+    content: string;
+}
+
 export interface ProjectInfoRequest {
     projectPath: string;
 }
@@ -2142,6 +2150,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>;

@@ -1436,6 +1436,7 @@
         "copyJSLibs": "copyfiles -f ../ballerina-visualizer/build/*.js resources/jslibs && copyfiles -f ../ballerina-visualizer/build/images/* resources/jslibs/images && copyfiles -f ../trace-visualizer/build/*.js resources/jslibs && copyfiles -f ../graphql/build/*.js resources/jslibs"
     },
     "dependencies": {
+        "@vscode/ripgrep": "1.15.9",
         "@ai-sdk/amazon-bedrock": "4.0.83",
         "@ai-sdk/anthropic": "3.0.64",
         "@ai-sdk/google-vertex": "4.0.94",

@@ -275,6 +275,8 @@ import {
     ClausePositionRequest,
     SemanticDiffRequest,
     SemanticDiffResponse,
+    CodeMapRequest,
+    CodeMapResponse,
     ConvertExpressionRequest,
     ConvertExpressionResponse,
     IntrospectDatabaseRequest,
@@ -491,6 +493,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',
@@ -1502,6 +1505,11 @@ 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 --------------->
 
 

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

@@ -270,7 +270,15 @@ export class AgentExecutor extends AICommandExecutor<GenerateAgentCodeRequest> {
             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);
@@ -296,7 +304,7 @@ 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",

@@ -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 || projectRootPath;
+        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;

@@ -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.
@@ -174,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.
+            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 <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}.
@@ -239,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) {

@@ -21,6 +21,8 @@ import { ExecutionContext, ProjectSource } from '@wso2/ballerina-core';
 import { CopilotEventHandler } from '../utils/events';
 import { createTaskWriteTool, TASK_WRITE_TOOL_NAME } from './tools/task-writer';
 import { createDiagnosticsTool, DIAGNOSTICS_TOOL_NAME } from './tools/diagnostics';
+import { createGrepTool, createGrepExecute, GREP_TOOL_NAME } from './tools/grep';
+import { createGlobTool, createGlobExecute, GLOB_TOOL_NAME } from './tools/glob';
 import {
     createBatchEditTool,
     createEditExecute,
@@ -122,6 +124,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),