Implement custom LS API for CodeMap generation in BI Copilot by yasithrashan · Pull Request #649 · ballerina-platform/ballerina-language-server

yasithrashan · 2026-01-21T09:45:14Z

Purpose

Related to: wso2/product-ballerina-integrator#2192

Ballerina Copilot currently sends the full project code to the LLM, which causes high token usage and slow responses.
To fix this, we generate a CodeMap of the project using a custom Language Server (LS) API and send this CodeMap to the LLM so it can understand and navigate the codebase without needing the full source.

Goals

Extract structured code artifacts from the project

Approach

Create a custom Ballerina Language Server (LS) API to generate a CodeMap for the project.
Use the API to extract artifacts from the entire project (services, functions, types, variables, listeners, etc.).

Copilot

Pull request overview

This PR implements a custom Language Server API to generate a CodeMap for Ballerina Copilot, reducing token usage by sending structured code artifacts instead of full source code to the LLM.

Changes:

Introduced CodeMap generation API with request/response models and service endpoint
Implemented AST transformation to extract structured artifacts (services, functions, types, listeners, connections, etc.)
Added comprehensive test coverage with multiple Ballerina project scenarios

Reviewed changes

Copilot reviewed 57 out of 57 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
CodeMapArtifact.java	Data model for code artifacts with builder pattern
CodeMapFile.java	Record representing a file with its artifacts
CodeMapGenerator.java	Main generator logic iterating through project modules
CodeMapNodeTransformer.java	AST visitor extracting artifacts from syntax nodes
CodeMapRequest.java	Request model with project path
CodeMapResponse.java	Response model with extracted files
DesignModelGeneratorService.java	Added codeMap endpoint to existing service
module-info.java	Exported codemap package
CodeMapGeneratorTest.java	Test class with data provider pattern
testng.xml	Added test class to test suite
codemap/source/*	Test Ballerina projects for various scenarios
codemap/config/*	Expected JSON output for test validation

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

...model-generator-ls-extension/src/test/resources/codemap/source/http_service/http_service.bal

...rchitecture-model-generator-ls-extension/src/test/resources/codemap/config/http_service.json

SasinduDilshara · 2026-01-21T15:04:08Z

...el-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapArtifact.java

+import java.util.List;
+import java.util.Map;
+
+public record CodeMapArtifact(String name, String type, LineRange lineRange, List<String> modifiers,


Children can be a optional field, give a default for that

SasinduDilshara · 2026-01-21T15:04:40Z

...el-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapArtifact.java

+import java.util.List;
+import java.util.Map;
+
+public record CodeMapArtifact(String name, String type, LineRange lineRange, List<String> modifiers,


Please give some samples for modifiers, propertirs, children. What are they?

SasinduDilshara · 2026-01-21T15:05:03Z

...el-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapArtifact.java

+import java.util.List;
+import java.util.Map;
+
+public record CodeMapArtifact(String name, String type, LineRange lineRange, List<String> modifiers,


Where is the filename capture in here?

The filename is not captured directly. In the CodeMapGenerator.java file, within the
generateCodeMap function, we get the file name using:

String fileName = document.name();

Then, in the same function, after extracting the artifacts from a single file, we wrap
everything into a CodeMapFile

CodeMapFile codeMapFile = new CodeMapFile(fileName, relativeFilePath, artifacts);

nipunayf · 2026-01-21T18:03:42Z

...-model-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapFile.java

+import java.util.Collections;
+import java.util.List;
+
+public record CodeMapFile(String fileName, String relativeFilePath, List<CodeMapArtifact> artifacts) {


Missing java doc.

Add java docs for all the public classes.

@yasithrashan This is mandatory.
All public constructs should have Javadocs while explaining the main functioanlity and inputs and outputs

And also please mark public only the necessary ones

nipunayf · 2026-01-25T13:49:34Z

...enerator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/ChangedFilesTracker.java

@@ -0,0 +1,68 @@
+/*
+ *  Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com)


Suggested change

* Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com)

* Copyright (c) 2026, WSO2 LLC. (http://www.wso2.com)

Apply this to other applicable areas.

nipunayf · 2026-01-25T13:50:04Z

...enerator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/ChangedFilesTracker.java

+ * Tracks changed files per project for incremental code map generation.
+ * This singleton maintains a thread-safe record of file changes between API calls.
+ *
+ * @since 1.0.0


Suggested change

* @since 1.0.0

* @since 1.6.0

Apply this to other applicable areas.

nipunayf · 2026-01-25T13:53:38Z

...enerator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/ChangedFilesTracker.java

+        if (instance == null) {
+            instance = new ChangedFilesTracker();
+        }
+        return instance;


Making the method synchronized may be overkill, since after the first call all subsequent calls are read-only. Use the holder method for handling the Singleton.

nipunayf · 2026-01-25T14:50:52Z

...enerator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/ChangedFilesTracker.java

+
+     // Get and clear the list of changed files for a project.
+    public List<String> getAndClearChangedFiles(String projectKey) {
+        Set<String> files = changedFilesMap.remove(projectKey);


We should implement a rollback mechanism for this instruction. If the LS fails to respond, the state becomes desynchronized for the extension because files are being removed prematurely. Given the low likelihood of this error, adding a note is sufficient for now.

nipunayf · 2026-01-25T17:56:07Z

...rator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapNodeTransformer.java

+        extractDocumentation(functionDefinitionNode.metadata()).ifPresent(functionBuilder::documentation);
+        extractComments(functionDefinitionNode).ifPresent(functionBuilder::comment);
+
+        functionBuilder.type("FUNCTION");


Extract the strings used to define the JSON fields to the class level.

nipunayf · 2026-01-25T17:56:52Z

...o/ballerina/architecturemodelgenerator/extension/ArchitectureModelGeneratorServiceTests.java

Do we need these changes for this PR?

nipunayf · 2026-01-25T17:57:03Z

...xtension/src/test/java/io/ballerina/designmodelgenerator/extension/CodeMapGeneratorTest.java

+
+        if (!files.equals(testConfig.output())) {
+            TestConfig updatedConfig = new TestConfig(testConfig.description(), testConfig.source(), files);
+            updateConfig(configJsonPath, updatedConfig);


Suggested change

updateConfig(configJsonPath, updatedConfig);

// updateConfig(configJsonPath, updatedConfig);

nipunayf · 2026-01-25T18:12:52Z

.../src/test/java/io/ballerina/designmodelgenerator/extension/PublishCodeMapSubscriberTest.java

+ *
+ * @since 1.0.0
+ */
+public class PublishCodeMapSubscriberTest extends AbstractLSTest {


Need to cover following scenarios if not included:

Multiple File Accumulation - Trigger onEvent for two different files within the same project sequentially.

Consecutive Events for Same File - Trigger onEvent multiple times for the same file.

Same File in Different Modules: Trigger events for both types.bal and modules/mod1/types.bal.

State Clearing - Verify that retrieving the changes "consumes" them, ensuring subsequent calls don't return stale data.

Project Isolation: Multiple Project - Verify that changes in "Project A" do not leak into "Project B".

nipunayf · 2026-01-25T18:16:56Z

...xtension/src/test/java/io/ballerina/designmodelgenerator/extension/CodeMapGeneratorTest.java

+ *
+ * @since 1.0.0
+ */
+public class CodeMapGeneratorTest extends AbstractLSTest {


Add a single test case with a Ballerina project as the source. Keep the code minimal, since the goal is only to verify that the API works for Ballerina projects, not to validate different construct types.

SasinduDilshara · 2026-01-26T05:09:13Z

...enerator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/ChangedFilesTracker.java

+
+    public static synchronized ChangedFilesTracker getInstance() {
+        if (instance == null) {
+            instance = new ChangedFilesTracker();


public static ChangedFilesTracker getInstance() { // First check (no locking) if (instance == null) { // Only synchronize if it looks like we need to create it synchronized (ChangedFilesTracker.class) { // Second check (inside the lock) if (instance == null) { instance = new ChangedFilesTracker(); } } } return instance; }

SasinduDilshara · 2026-01-26T05:11:46Z

...enerator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapFilesTracker.java

+    public List<String> getAndClearChangedFiles(String projectKey) {
+        Set<String> files = changedFilesMap.remove(projectKey);
+        if (files == null || files.isEmpty()) {
+            return Collections.emptyList();


Lets add a log for this

SasinduDilshara · 2026-01-26T05:12:21Z

...el-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapArtifact.java

+
+        public Builder modifiers(List<String> modifiers) {
+            if (!modifiers.isEmpty()) {
+                this.properties.put("modifiers", new ArrayList<>(modifiers));


This kind of literals should be defines as constants

SasinduDilshara · 2026-01-26T05:12:51Z

...el-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapArtifact.java

+        }
+
+        public Builder config(String config) {
+            return addProperty("config", config);


Is the config is a simple element?

SasinduDilshara · 2026-01-26T05:15:25Z

...l-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapGenerator.java

+            ModuleInfo moduleInfo = createModuleInfo(module.descriptor());
+
+            // Iterate through each document in the module
+            for (var documentId : module.documentIds()) {


This may not contains resources

In here what I mean by resources is,
In a Ballerina project there can be a folder called resources.

It may contains images and other necessaru resources

Basically its not a bal file, But images and other resources necessaru for the code

SasinduDilshara · 2026-01-26T05:16:24Z

...l-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapGenerator.java

+
+        for (var moduleId : currentPackage.moduleIds()) {
+            Module module = currentPackage.module(moduleId);
+            SemanticModel semanticModel =


Handle scearios where the semanticModel = null

SasinduDilshara · 2026-01-26T05:17:45Z

...l-generator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapGenerator.java

+                    PackageUtil.getCompilation(currentPackage).getSemanticModel(moduleId);
+            ModuleInfo moduleInfo = createModuleInfo(module.descriptor());
+
+            for (var documentId : module.documentIds()) {


Add a variable isTest = true/false, Depends on tat add test files also into here. It will be helpful when copilot generating tests

SasinduDilshara · 2026-01-26T05:20:34Z

...rator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapNodeTransformer.java

+    private final ModuleInfo moduleInfo;
+    private final boolean extractComments;
+
+    private static final String AUTOMATION_FUNCTION_NAME = "automation";


Please check, there is no syntax called automation

SasinduDilshara · 2026-01-26T05:25:42Z

...rator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapNodeTransformer.java

+        // Build full import name
+        String fullImportName = orgName.isEmpty() ? moduleName : orgName + "/" + moduleName;
+        if (alias.isPresent()) {
+            fullImportName += " as " + alias.get();


Add this terms as constnats.

SasinduDilshara · 2026-01-26T05:27:16Z

...rator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapNodeTransformer.java

+        } else if (firstExpression != null) {
+            return firstExpression.toSourceCode().strip();
+        } else {
+            return "";


There is no meaning of empty string as service name

Use Optional.empty

SasinduDilshara · 2026-01-26T05:27:49Z

...rator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapNodeTransformer.java

+                return Optional.of(classSymbol);
+            }
+        } catch (Throwable e) {
+            // Ignore


SasinduDilshara · 2026-01-26T05:28:20Z

...rator-core/src/main/java/io/ballerina/artifactsgenerator/codemap/CodeMapNodeTransformer.java

+                .filter(doc -> !doc.isEmpty());
+    }
+
+    private Optional<String> extractComments(Node node) {


Suggested change

private Optional<String> extractComments(Node node) {

private Optional<String> extractInlineComments(Node node) {

SasinduDilshara · 2026-01-26T05:28:42Z

...n/src/main/java/io/ballerina/designmodelgenerator/extension/DesignModelGeneratorService.java

+                        .getAndClearChangedFiles(projectKey);
+
+                if (changedFiles.isEmpty()) {
+                    // No changes tracked, return empty response


SasinduDilshara · 2026-01-26T05:30:11Z

...s/architecture-model-generator-ls-extension/src/test/resources/codemap/source/connection.bal

+import ballerina/tcp;
+
+// HTTP connection
+final http:Client httpConnection = check new ("https://api.example.com", {


Add test for Url as varaibales as well

coderabbitai · 2026-02-09T05:42:12Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

yasithrashan force-pushed the feature/custom-ls-api branch from 9aead0c to caa5ccd Compare January 21, 2026 09:55

nipunayf requested a review from Copilot January 21, 2026 10:47

Copilot started reviewing on behalf of nipunayf January 21, 2026 10:48 View session

Copilot AI reviewed Jan 21, 2026

View reviewed changes

...model-generator-ls-extension/src/test/resources/codemap/source/http_service/http_service.bal Outdated Show resolved Hide resolved

...rchitecture-model-generator-ls-extension/src/test/resources/codemap/config/http_service.json Outdated Show resolved Hide resolved

SasinduDilshara reviewed Jan 21, 2026

View reviewed changes

yasithrashan force-pushed the feature/custom-ls-api branch from 254bb0e to 2b2b5bc Compare January 23, 2026 09:45

nipunayf reviewed Jan 25, 2026

View reviewed changes

SasinduDilshara reviewed Jan 26, 2026

View reviewed changes

yasithrashan force-pushed the feature/custom-ls-api branch from 1f255a3 to 0d9f588 Compare January 27, 2026 08:12

yasithrashan force-pushed the feature/custom-ls-api branch from 0d9f588 to 0304b48 Compare February 5, 2026 09:00

yasithrashan added 4 commits February 9, 2026 11:12

Add CodeMap artifact generation core

65fdb0b

Add code map generation and LS extension support

e85f9bb

Add CodeMap test resources and configs

9073498

Add class field support, improve types

9c93f35

yasithrashan added 24 commits February 9, 2026 11:12

Add listener type extraction to service

cf099aa

Improve type and record handling in CodeMapNodeTransformer

dae5b84

Add listener init args and improve record types

db1df80

Add relative path and submodule support

326181f

Add documentation extraction to CodeMap artifacts

1712428

Add method to generate code map for specific files

2893709

Add comment extraction to CodeMapNodeTransformer

5dfcf74

Fix Checkstyle violations

dfef6e2

Add incremental code map generation support

b247bf3

Use system line separator for joining comments

c4b3b06

Refactor test resources and update config references

128cc9c

Move modifiers into properties in CodeMapArtifact

55e096f

Refactor artifact type handling and add category property

cd114fe

Add import and constant support to CodeMap generator

c4c2325

Merge codeMap APIs and support changed files

7f991e5

Refactor code map API and improve artifact structure

b333124

Refactor CodeMap generator and tests

8737b4a

Refactor property keys and comment extraction methods

45a92d1

Remove unused lineRange method from Builder

af0c1f4

Update copyright year and @SInCE tags, fix test paths

6be916b

Track changed files by relative path instead of name

7f93769

Track didOpen events for CodeMap files

52561b0

Fix main function name in CodeMap

f2a56c9

Update CodeMap to remove redundant fields

e760d4c

yasithrashan force-pushed the feature/custom-ls-api branch from c486070 to e760d4c Compare February 9, 2026 05:43

yasithrashan added 5 commits February 9, 2026 11:43

Rename lineRange to range and drop filePath

37f8aa1

Rename AUTOMATION category to MAIN

11b6657

Remove category field from code map artifacts

03da945

Always create CodeMapFile and reorder imports

73a9ac4

Add enum type handling to CodeMap transformer

56619ae

		@@ -0,0 +1,68 @@
		/*
		* Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com)

	* Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com)
	* Copyright (c) 2026, WSO2 LLC. (http://www.wso2.com)

	updateConfig(configJsonPath, updatedConfig);
	// updateConfig(configJsonPath, updatedConfig);

	private Optional<String> extractComments(Node node) {
	private Optional<String> extractInlineComments(Node node) {

Comments

Conversation

yasithrashan commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Goals

Approach

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SasinduDilshara Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nipunayf Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SasinduDilshara Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

yasithrashan commented Jan 21, 2026 •

edited

Loading

SasinduDilshara Jan 21, 2026 •

edited

Loading

nipunayf Jan 21, 2026 •

edited

Loading

SasinduDilshara Jan 26, 2026 •

edited

Loading