From c4f26d1c0e1b1362228c15af967f6454ec4c9404 Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Wed, 21 Jan 2026 15:07:13 +0530 Subject: [PATCH 1/8] Make anon http responses editable --- .../extension/model/HttpResponse.java | 19 +++++++- .../extension/util/HttpUtil.java | 46 ++++++++++--------- 2 files changed, 41 insertions(+), 24 deletions(-) diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/model/HttpResponse.java b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/model/HttpResponse.java index 7e59e8d24b..10cc7129c5 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/model/HttpResponse.java +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/model/HttpResponse.java @@ -37,13 +37,16 @@ public class HttpResponse { private boolean editable = false; private boolean advanced = false; - public HttpResponse(Value statusCode, Value body, Value mediaType, Value name, Value type, Value headers) { + public HttpResponse(Value statusCode, Value body, Value mediaType, Value name, Value type, Value headers, + boolean editable, boolean enabled) { this.statusCode = statusCode; this.body = body; this.mediaType = mediaType; this.name = name; this.type = type; this.headers = headers; + this.editable = editable; + this.enabled = enabled; } public HttpResponse(String type) { @@ -156,6 +159,8 @@ public static class Builder { private Value name; private Value type; private Value headers; + private boolean editable = false; + private boolean enabled = false; public Builder statusCode(String statusCode, boolean editable) { this.statusCode = createValue(statusCode, Value.FieldType.SINGLE_SELECT, editable); @@ -188,6 +193,16 @@ public Builder headers(Object headers, boolean editable) { return this; } + public Builder editable(boolean editable) { + this.editable = editable; + return this; + } + + public Builder enabled(boolean enabled) { + this.enabled = enabled; + return this; + } + public HttpResponse build() { if (mediaType == null) { this.mediaType = createValue("", Value.FieldType.EXPRESSION, true); @@ -204,7 +219,7 @@ public HttpResponse build() { if (body == null) { this.body = createOptionalValue("", Value.FieldType.EXPRESSION, true); } - return new HttpResponse(statusCode, body, mediaType, name, type, headers); + return new HttpResponse(statusCode, body, mediaType, name, type, headers, editable, enabled); } } } diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/HttpUtil.java b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/HttpUtil.java index a31dd36b51..a4b3676f1f 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/HttpUtil.java +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/HttpUtil.java @@ -429,13 +429,13 @@ private static List getHttpResponses(TypeSymbol returnTypeSymbol, TypeSymbol xmlType = typeBuilder.XML_TYPE.build(); anydataResponses.forEach(type -> { - HttpResponse.Builder builder = new HttpResponse.Builder() + HttpResponse response = new HttpResponse.Builder() .statusCode(String.valueOf(defaultStatusCode), true) .body(getTypeName(type, currentModuleName), true) - .mediaType(deriveMediaType(type, stringType, byteArrayType, xmlType), true); - HttpResponse response = builder.build(); - response.setEnabled(true); - response.setEditable(true); + .mediaType(deriveMediaType(type, stringType, byteArrayType, xmlType), true) + .editable(true) + .enabled(true) + .build(); responses.add(response); }); @@ -608,33 +608,34 @@ private static HttpResponse getHttpResponse(TypeSymbol statusCodeResponseType, S String statusCode = getResponseCode(statusCodeResponseType, defaultStatusCode, semanticModel); String signature = statusCodeResponseType.signature().trim(); if (signature.startsWith("record {") && signature.endsWith("}")) { - HttpResponse httpResponse = buildHttpResponseFromTypeSymbol(statusCodeResponseType, currentModuleName, - statusCode, null); - httpResponse.setEditable(true); - return httpResponse; + return buildHttpResponseFromTypeSymbol(statusCodeResponseType, currentModuleName, statusCode, + null, true); } String typeName = getTypeName(statusCodeResponseType, currentModuleName); if (typeName.startsWith("http:")) { String type = HTTP_CODES_DES.get(statusCode); if (Objects.nonNull(type) && "http:%s".formatted(type).equals(typeName)) { - HttpResponse.Builder builder = new HttpResponse.Builder() + return new HttpResponse.Builder() .statusCode(statusCode, true) .type(typeName, true) .body("", true) .name("", true) .headers("", true) - .mediaType("", true); - return builder.build(); + .mediaType("", true) + .enabled(true) + .editable(true) + .build(); } } - return buildHttpResponseFromTypeSymbol(statusCodeResponseType, currentModuleName, statusCode, typeName); + return buildHttpResponseFromTypeSymbol(statusCodeResponseType, currentModuleName, statusCode, typeName, false); } private static HttpResponse buildHttpResponseFromTypeSymbol(TypeSymbol statusCodeResponseType, String currentModuleName, String statusCode, - String typeName) { + String typeName, + boolean editable) { List headers = new ArrayList<>(); String body = "anydata"; String mediaType = ""; @@ -659,14 +660,15 @@ private static HttpResponse buildHttpResponseFromTypeSymbol(TypeSymbol statusCod } } } - HttpResponse.Builder builder = new HttpResponse.Builder() - .statusCode(statusCode, false) - .type(typeName, false) - .body(body, false) - .headers(headers, false) - .name("", false) - .mediaType(mediaType, false); - return builder.build(); + return new HttpResponse.Builder() + .statusCode(statusCode, editable) + .type(typeName, editable) + .body(body, editable) + .headers(headers, editable) + .name("", editable) + .mediaType(mediaType, editable) + .editable(editable) + .build(); } public static boolean isSubTypeOfHttpStatusCodeResponse(TypeSymbol typeSymbol, SemanticModel semanticModel) { From 5991a81a397905ed367538949e84be58fc67f040 Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Wed, 21 Jan 2026 15:09:00 +0530 Subject: [PATCH 2/8] Update get function model from source tests --- .../config/http_methods_delete.json | 4 ++-- .../config/http_methods_head.json | 2 +- .../config/http_methods_options.json | 2 +- .../config/http_methods_post.json | 2 +- .../config/http_methods_put.json | 2 +- .../config/http_response_2.json | 2 +- .../config/http_response_4.json | 10 +++++----- .../config/http_response_5.json | 18 +++++++++--------- 8 files changed, 21 insertions(+), 21 deletions(-) diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_delete.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_delete.json index 8a392475f6..e9d14db4fb 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_delete.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_delete.json @@ -477,7 +477,7 @@ "advanced": false }, "enabled": true, - "editable": false, + "editable": true, "advanced": false }, { @@ -566,7 +566,7 @@ "advanced": false }, "enabled": true, - "editable": false, + "editable": true, "advanced": false } ], diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_head.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_head.json index 7c1a974271..e694b6bc1e 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_head.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_head.json @@ -477,7 +477,7 @@ "advanced": false }, "enabled": true, - "editable": false, + "editable": true, "advanced": false } ], diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_options.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_options.json index c67d89febd..547d7e7da0 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_options.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_options.json @@ -477,7 +477,7 @@ "advanced": false }, "enabled": true, - "editable": false, + "editable": true, "advanced": false } ], diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_post.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_post.json index 89409c5b97..1694f625af 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_post.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_post.json @@ -570,7 +570,7 @@ "advanced": false }, "enabled": true, - "editable": false, + "editable": true, "advanced": false }, { diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_put.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_put.json index 134ebda03e..7a439776f0 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_put.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_methods_put.json @@ -659,7 +659,7 @@ "advanced": false }, "enabled": true, - "editable": false, + "editable": true, "advanced": false } ], diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_2.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_2.json index 4077f69989..9dbd22ab00 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_2.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_2.json @@ -477,7 +477,7 @@ "advanced": false }, "enabled": true, - "editable": false, + "editable": true, "advanced": false } ], diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_4.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_4.json index 7798e01e02..61b2c82661 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_4.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_4.json @@ -476,7 +476,7 @@ "optional": true, "advanced": false }, - "enabled": true, + "enabled": false, "editable": false, "advanced": false }, @@ -565,7 +565,7 @@ "optional": true, "advanced": false }, - "enabled": true, + "enabled": false, "editable": false, "advanced": false }, @@ -645,8 +645,8 @@ "value": [ { "type": "string", - "optional": false, - "name": "x\\-path" + "name": "x\\-path", + "optional": false } ], "types": [ @@ -660,7 +660,7 @@ "optional": true, "advanced": false }, - "enabled": true, + "enabled": false, "editable": false, "advanced": false } diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_5.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_5.json index e7171b4564..eab5d26a41 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_5.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/get_fm_from_source/config/http_response_5.json @@ -402,7 +402,7 @@ } ], "enabled": true, - "editable": false, + "editable": true, "optional": false, "advanced": false }, @@ -416,7 +416,7 @@ } ], "enabled": true, - "editable": false, + "editable": true, "optional": true, "advanced": false }, @@ -430,7 +430,7 @@ } ], "enabled": true, - "editable": false, + "editable": true, "optional": false, "advanced": false }, @@ -444,7 +444,7 @@ } ], "enabled": true, - "editable": false, + "editable": true, "optional": true, "advanced": false }, @@ -457,7 +457,7 @@ } ], "enabled": true, - "editable": false, + "editable": true, "optional": true, "advanced": false }, @@ -466,8 +466,8 @@ "value": [ { "type": "string", - "optional": false, - "name": "x\\-path" + "name": "x\\-path", + "optional": false } ], "types": [ @@ -477,11 +477,11 @@ } ], "enabled": true, - "editable": false, + "editable": true, "optional": true, "advanced": false }, - "enabled": true, + "enabled": false, "editable": true, "advanced": false } From 3af1eae702c56c44df8bca48b95e269f0e00fa18 Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Wed, 21 Jan 2026 17:33:30 +0530 Subject: [PATCH 3/8] Add tests for adding kafka func --- .../config/add_kafka_on_consumer_record.json | 309 ++++++++++++++++++ .../source/sample1/Ballerina.toml | 6 + .../add_function/source/sample1/main.bal | 6 + .../add_function/source/sample1/types.bal | 0 4 files changed, 321 insertions(+) create mode 100644 service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json create mode 100644 service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/Ballerina.toml create mode 100644 service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/main.bal create mode 100644 service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/types.bal diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json new file mode 100644 index 0000000000..705c50bbcd --- /dev/null +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json @@ -0,0 +1,309 @@ +{ + "filePath": "sample1/main.bal", + "description": "Add onConsumerRecord remote method to the Kafka service", + "codedata": { + "lineRange": { + "startLine": { + "line": 4, + "offset": 0 + }, + "endLine": { + "line": 5, + "offset": 1 + } + } + }, + "function": { + "metadata": { + "label": "onConsumerRecord", + "description": "The `onConsumerRecord` remote method will be triggered when a message is received from Kafka topic(s)" + }, + "kind": "REMOTE", + "name": { + "metadata": { + "label": "onConsumerRecord", + "description": "The `onConsumerRecord` remote method will be triggered when a message is received from Kafka topic(s)" + }, + "placeholder": "onConsumerRecord", + "value": "onConsumerRecord", + "types": [ + { + "fieldType": "IDENTIFIER", + "selected": false + } + ], + "enabled": true, + "editable": false, + "optional": false, + "advanced": false + }, + "documentation": { + "types": [ + { + "fieldType": "TEXT", + "selected": true + } + ], + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "parameters": [ + { + "metadata": { + "label": "messages", + "description": "The messages received for Kafka topic(s)" + }, + "kind": "REQUIRED", + "type": { + "metadata": { + "label": "Parameter Type", + "description": "The type of the parameter" + }, + "placeholder": "kafka:AnydataConsumerRecord[]", + "value": "kafka:AnydataConsumerRecord[]", + "types": [ + { + "fieldType": "TYPE", + "selected": true + } + ], + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "name": { + "metadata": { + "label": "messages", + "description": "The messages received for Kafka topic(s)" + }, + "placeholder": "messages", + "value": "messages", + "types": [ + { + "fieldType": "IDENTIFIER", + "selected": true + } + ], + "enabled": true, + "editable": true, + "optional": false, + "advanced": false + }, + "defaultValue": { + "metadata": { + "label": "Default Value", + "description": "The default value of the parameter" + }, + "placeholder": "", + "value": "", + "types": [ + { + "fieldType": "EXPRESSION", + "selected": true + } + ], + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "enabled": true, + "editable": true, + "optional": false, + "advanced": false, + "hidden": false, + "isGraphqlId": false + }, + { + "metadata": { + "label": "caller", + "description": "The Kafka caller object to commit the offsets" + }, + "kind": "OPTIONAL", + "type": { + "metadata": { + "label": "Parameter Type", + "description": "The type of the parameter" + }, + "placeholder": "kafka:Caller", + "value": "kafka:Caller", + "types": [ + { + "fieldType": "TYPE", + "selected": true + } + ], + "enabled": true, + "editable": false, + "optional": true, + "advanced": false + }, + "name": { + "metadata": { + "label": "caller", + "description": "The Kafka caller object to commit the offsets" + }, + "placeholder": "caller", + "value": "caller", + "types": [ + { + "fieldType": "IDENTIFIER", + "selected": true + } + ], + "enabled": true, + "editable": true, + "optional": false, + "advanced": false + }, + "defaultValue": { + "metadata": { + "label": "Default Value", + "description": "The default value of the parameter" + }, + "placeholder": "", + "value": "", + "types": [ + { + "fieldType": "EXPRESSION", + "selected": true + } + ], + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "enabled": true, + "editable": true, + "optional": true, + "advanced": false, + "hidden": false, + "isGraphqlId": false + }, + { + "metadata": { + "label": "Data Binding", + "description": "Data binding parameter" + }, + "kind": "DATA_BINDING", + "type": { + "value": "ai:Schema", + "types": [ + { + "fieldType": "TYPE", + "selected": true + } + ], + "enabled": true, + "editable": false, + "optional": false, + "advanced": false, + "imports": { + "ai": "ballerina/ai:1.6.1" + } + }, + "name": { + "value": "payload", + "types": [ + { + "fieldType": "IDENTIFIER", + "selected": true + } + ], + "enabled": true, + "editable": false, + "optional": false, + "advanced": false + }, + "enabled": true, + "editable": true, + "optional": false, + "advanced": false, + "hidden": false, + "isGraphqlId": false + } + ], + "returnType": { + "hasError": true, + "isGraphqlId": false, + "metadata": { + "label": "Return Type", + "description": "The return type of the function" + }, + "placeholder": "error?", + "value": "error?", + "types": [ + { + "fieldType": "TYPE", + "selected": true + } + ], + "enabled": true, + "editable": false, + "optional": true, + "advanced": false + }, + "enabled": true, + "optional": false, + "editable": false, + "canAddParameters": false, + "codedata": { + "lineRange": { + "fileName": "main.bal", + "startLine": { + "line": 6, + "offset": 4 + }, + "endLine": { + "line": 11, + "offset": 5 + } + }, + "moduleName": "kafka" + }, + "properties": { + "wrapperTypeName": { + "value": "KafkaAnydataConsumer", + "enabled": false, + "editable": false, + "optional": false, + "advanced": false + }, + "payloadFieldName": { + "value": "value", + "enabled": false, + "editable": false, + "optional": false, + "advanced": false + }, + "canDataBind": { + "value": "true", + "enabled": false, + "editable": false, + "optional": false, + "advanced": false + } + } + }, + "output": { + "sample1/main.bal": [ + { + "range": { + "start": { + "line": 4, + "character": 40 + }, + "end": { + "line": 4, + "character": 40 + } + }, + "newText": "\n\tremote function onConsumerRecord(KafkaAnydataConsumer[] messages, kafka:Caller caller) returns error? {\n\t\tdo {\n\t\t} on fail error err {\n\t\t\t// handle error\n\t\t\treturn error(\"unhandled error\", err);\n\t\t}\n\t}" + } + ] + } +} diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/Ballerina.toml b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/Ballerina.toml new file mode 100644 index 0000000000..4c600c3a38 --- /dev/null +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/Ballerina.toml @@ -0,0 +1,6 @@ +[package] +org = "ballerina" +name = "service_designer" +version = "0.1.0" + +bi = true diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/main.bal b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/main.bal new file mode 100644 index 0000000000..9bd7424dc0 --- /dev/null +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/main.bal @@ -0,0 +1,6 @@ +import ballerinax/kafka; + +listener kafka:Listener kafkaListener = new ("localhost:9092", topics = "myTopic"); + +service kafka:Service on kafkaListener { +} diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/types.bal b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/source/sample1/types.bal new file mode 100644 index 0000000000..e69de29bb2 From 7d765c7a05c9cc290200d0c2db1bc5eaa81bf7e9 Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Wed, 21 Jan 2026 17:45:51 +0530 Subject: [PATCH 4/8] Support data binding for imported types --- .../extension/util/DatabindUtil.java | 53 ++++++++++++++----- 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java index 745c3b220c..7a9ddaf751 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java @@ -513,8 +513,10 @@ public static Map> processDatabindingForAdd(AddModelConte new Value.ValueBuilder().value(typeName).build() ); - return createTypeDefinitionEdits(context.project(), typeName, baseType, - newDataBindingType, payloadFieldName, context.filePath(), context.workspaceManager()); + Map importsForTypeDef = dataBindingParam.getType().getImports(); + + return createTypeDefinitionEdits(context.project(), typeName, baseType, newDataBindingType, payloadFieldName, + context.filePath(), context.workspaceManager(), importsForTypeDef); } /** @@ -625,27 +627,29 @@ public static Map> processDatabindingUpdate(UpdateModelCo Map> typesEdits; String typeName; + Map importsForTypeDef = dataBindingParam.getType().getImports(); if (customWrapperTypeName != null && !customWrapperTypeName.equals(existingTypeName)) { typeName = customWrapperTypeName; if (existingTypeName != null) { typesEdits = updateTypeDefinitionEdits(context, existingTypeName, baseType, newDataBindingType, - payloadFieldName, customWrapperTypeName); + payloadFieldName, customWrapperTypeName, importsForTypeDef); } else { typesEdits = createTypeDefinitionEdits(context.project(), customWrapperTypeName, baseType, - newDataBindingType, payloadFieldName, context.filePath(), context.workspaceManager()); + newDataBindingType, payloadFieldName, context.filePath(), context.workspaceManager(), + importsForTypeDef); } } else if (existingTypeName != null) { typeName = existingTypeName; typesEdits = updateTypeDefinitionEdits(context, existingTypeName, baseType, newDataBindingType, - payloadFieldName, null); + payloadFieldName, null, importsForTypeDef); } else { typeName = generateNewDataBindTypeName(context.filePath(), context.workspaceManager(), context.semanticModel(), context.functionNode(), prefix); typesEdits = createTypeDefinitionEdits(context.project(), typeName, baseType, newDataBindingType, - payloadFieldName, context.filePath(), context.workspaceManager()); + payloadFieldName, context.filePath(), context.workspaceManager(), importsForTypeDef); } updateFunctionParameters(function, dataBindingParam, typeName, isArray); @@ -885,9 +889,11 @@ private static String generateNewDataBindTypeName(String contextFilePath, Worksp * * @param baseType The base record type (e.g., "kafka:AnydataConsumerRecord") * @param modulePartNode The module part node to check existing imports + * @param importsForTypeDef Map of imports needed for the type definition * @return Set of import statements to add */ - private static Set extractRequiredImports(String baseType, ModulePartNode modulePartNode) { + private static Set extractRequiredImports(String baseType, ModulePartNode modulePartNode, + Map importsForTypeDef) { Set imports = new HashSet<>(); if (baseType.contains(COLON)) { @@ -900,6 +906,15 @@ private static Set extractRequiredImports(String baseType, ModulePartNod } } + importsForTypeDef.values().forEach(moduleId -> { + String[] importParts = moduleId.split("/"); + String orgName = importParts[0]; + String moduleName = importParts[1].split(":")[0]; + if (!importExists(modulePartNode, orgName, moduleName)) { + imports.add(getImportStmt(orgName, moduleName)); + } + }); + return imports; } @@ -915,14 +930,15 @@ private static Set extractRequiredImports(String baseType, ModulePartNod */ private static TypeDefinitionEditContext prepareTypeDefinitionEditContext(Project project, String baseType, String contextFilePath, - WorkspaceManager workspaceManager) { + WorkspaceManager workspaceManager, + Map importsForTypeDef) { Document typesDocument = getTypesDocument(contextFilePath, workspaceManager); if (typesDocument == null || typesDocument.syntaxTree() == null) { return null; } ModulePartNode modulePartNode = typesDocument.syntaxTree().rootNode(); - Set requiredImports = extractRequiredImports(baseType, modulePartNode); + Set requiredImports = extractRequiredImports(baseType, modulePartNode, importsForTypeDef); return new TypeDefinitionEditContext(typesDocument, modulePartNode, requiredImports, project); } @@ -935,15 +951,20 @@ private static TypeDefinitionEditContext prepareTypeDefinitionEditContext(Projec * @param baseType The base record type (e.g., "kafka:AnydataConsumerRecord") * @param dataBindingType The data binding field type (e.g., "Order") * @param payloadFieldName The field name for the payload (e.g., "value" or "content") + * @param contextFilePath The context file path for locating types.bal + * @param workspaceManager The workspace manager for document retrieval + * @param importsForTypeDef Map of imports needed for the type definition * @return Map of file paths to TextEdit lists */ private static Map> createTypeDefinitionEdits(Project project, String typeName, String baseType, String dataBindingType, String payloadFieldName, String contextFilePath, - WorkspaceManager workspaceManager) { + WorkspaceManager workspaceManager, + Map importsForTypeDef) { TypeDefinitionEditContext context = - prepareTypeDefinitionEditContext(project, baseType, contextFilePath, workspaceManager); + prepareTypeDefinitionEditContext(project, baseType, contextFilePath, workspaceManager, + importsForTypeDef); if (context == null) { return Map.of(); } @@ -1172,6 +1193,7 @@ private static Path getFilePathForFile(String contextFilePath, WorkspaceManager * @param newDataBindingType The new data binding field type (e.g., "Customer") * @param payloadFieldName The field name for the payload (e.g., "value") * @param newTypeName The new type name to rename to (optional, if null uses existingTypeName) + * @param importsForTypeDef Map of imports needed for the type definition * @return Map of file paths to TextEdit lists */ private static Map> updateTypeDefinitionEdits(UpdateModelContext context, @@ -1179,10 +1201,13 @@ private static Map> updateTypeDefinitionEdits(UpdateModel String baseType, String newDataBindingType, String payloadFieldName, - String newTypeName) { + String newTypeName, + Map importsForTypeDef) { + Project project = context.project() != null ? context.project() : context.document().module().project(); TypeDefinitionEditContext editContext = - prepareTypeDefinitionEditContext(project, baseType, context.filePath(), context.workspaceManager()); + prepareTypeDefinitionEditContext(project, baseType, context.filePath(), context.workspaceManager(), + importsForTypeDef); if (editContext == null) { return Map.of(); } @@ -1203,7 +1228,7 @@ private static Map> updateTypeDefinitionEdits(UpdateModel if (existingTypeDef == null) { // Type doesn't exist, create it instead return createTypeDefinitionEdits(context.project(), existingTypeName, baseType, newDataBindingType, - payloadFieldName, context.filePath(), context.workspaceManager()); + payloadFieldName, context.filePath(), context.workspaceManager(), importsForTypeDef); } // Use newTypeName if provided for renaming, otherwise use existingTypeName From 714ff7666a4ac76802d33278692b5dd8c980437d Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Wed, 21 Jan 2026 17:47:49 +0530 Subject: [PATCH 5/8] Update tests --- .../config/add_kafka_on_consumer_record.json | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json index 705c50bbcd..7274cf6c90 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/add_function/config/add_kafka_on_consumer_record.json @@ -304,6 +304,34 @@ }, "newText": "\n\tremote function onConsumerRecord(KafkaAnydataConsumer[] messages, kafka:Caller caller) returns error? {\n\t\tdo {\n\t\t} on fail error err {\n\t\t\t// handle error\n\t\t\treturn error(\"unhandled error\", err);\n\t\t}\n\t}" } + ], + "sample1/types.bal": [ + { + "range": { + "start": { + "line": 0, + "character": 0 + }, + "end": { + "line": 0, + "character": 0 + } + }, + "newText": "\nimport ballerinax/kafka;\n\n\nimport ballerina/ai;\n\n" + }, + { + "range": { + "start": { + "line": 0, + "character": 0 + }, + "end": { + "line": 0, + "character": 0 + } + }, + "newText": "type KafkaAnydataConsumer record {|\n *kafka:AnydataConsumerRecord;\n ai:Schema value;\n|};" + } ] } } From c6483e18c89aa33b9b6e6793acbc82641012aa22 Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Wed, 21 Jan 2026 18:11:47 +0530 Subject: [PATCH 6/8] Add tests for updating kafka func --- .../config/kafka_bind_imported_type.json | 301 ++++++++++++++++++ .../source/kafka_databind/Ballerina.toml | 4 + 2 files changed, 305 insertions(+) create mode 100644 service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/config/kafka_bind_imported_type.json create mode 100644 service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/source/kafka_databind/Ballerina.toml diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/config/kafka_bind_imported_type.json b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/config/kafka_bind_imported_type.json new file mode 100644 index 0000000000..62f82551a0 --- /dev/null +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/config/kafka_bind_imported_type.json @@ -0,0 +1,301 @@ +{ + "filePath": "kafka_databind/main.bal", + "description": "Remove data binding from kafka event handler", + "function": { + "metadata": { + "label": "onConsumerRecord", + "description": "The `onConsumerRecord` remote method will be triggered when a message is received from Kafka topic(s)" + }, + "kind": "REMOTE", + "name": { + "metadata": { + "label": "onConsumerRecord", + "description": "The `onConsumerRecord` remote method will be triggered when a message is received from Kafka topic(s)" + }, + "codedata": { + "type": "FUNCTION_NAME" + }, + "placeholder": "onConsumerRecord", + "valueType": "IDENTIFIER", + "valueTypeConstraint": "string", + "value": "onConsumerRecord", + "enabled": true, + "editable": false, + "optional": false, + "advanced": false + }, + "documentation": { + "valueType": "STRING", + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "parameters": [ + { + "metadata": { + "label": "Data Binding", + "description": "Data binding parameter" + }, + "kind": "DATA_BINDING", + "type": { + "valueType": "TYPE", + "value": "ai:Schema", + "enabled": true, + "editable": false, + "optional": false, + "advanced": false, + "imports": { + "ai": "ballerina/ai:1.6.1" + } + }, + "name": { + "valueType": "IDENTIFIER", + "value": "records", + "enabled": true, + "editable": false, + "optional": false, + "advanced": false + }, + "documentation": { + "valueType": "STRING", + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "enabled": true, + "editable": true, + "optional": false, + "advanced": false, + "hidden": false, + "isGraphqlId": false + }, + { + "metadata": { + "label": "messages", + "description": "The messages received for Kafka topic(s)" + }, + "kind": "REQUIRED", + "type": { + "metadata": { + "label": "Type", + "description": "The type of the parameter" + }, + "placeholder": "kafka:AnydataConsumerRecord[]", + "valueType": "TYPE", + "value": "kafka:AnydataConsumerRecord[]", + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "name": { + "metadata": { + "label": "messages", + "description": "The messages received for Kafka topic(s)" + }, + "codedata": { + "type": "PARAMETER_NAME" + }, + "placeholder": "messages", + "valueType": "IDENTIFIER", + "value": "messages", + "enabled": true, + "editable": true, + "optional": false, + "advanced": false + }, + "defaultValue": { + "metadata": { + "label": "Default Value", + "description": "The default value of the parameter" + }, + "placeholder": "", + "valueType": "EXPRESSION", + "value": "", + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "documentation": { + "valueType": "STRING", + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "enabled": false, + "editable": true, + "optional": false, + "advanced": false, + "hidden": false, + "isGraphqlId": false + }, + { + "metadata": { + "label": "caller", + "description": "The Kafka caller object to commit the offsets" + }, + "kind": "OPTIONAL", + "type": { + "metadata": { + "label": "Type", + "description": "The type of the parameter" + }, + "placeholder": "kafka:Caller", + "valueType": "TYPE", + "value": "kafka:Caller", + "enabled": true, + "editable": false, + "optional": true, + "advanced": false + }, + "name": { + "metadata": { + "label": "caller", + "description": "The Kafka caller object to commit the offsets" + }, + "codedata": { + "type": "PARAMETER_NAME" + }, + "placeholder": "caller", + "valueType": "IDENTIFIER", + "value": "caller", + "enabled": true, + "editable": true, + "optional": false, + "advanced": false + }, + "defaultValue": { + "metadata": { + "label": "Default Value", + "description": "The default value of the parameter" + }, + "placeholder": "", + "valueType": "EXPRESSION", + "value": "", + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "documentation": { + "valueType": "STRING", + "enabled": true, + "editable": true, + "optional": true, + "advanced": false + }, + "enabled": false, + "editable": true, + "optional": true, + "advanced": false, + "hidden": false, + "isGraphqlId": false + } + ], + "returnType": { + "hasError": true, + "isGraphqlId": false, + "metadata": { + "label": "Return Type", + "description": "The return type of the function" + }, + "placeholder": "error?", + "valueType": "TYPE", + "value": "error?", + "enabled": true, + "editable": false, + "optional": true, + "advanced": false + }, + "enabled": true, + "optional": false, + "editable": false, + "canAddParameters": false, + "codedata": { + "lineRange": { + "fileName": "main.bal", + "startLine": { + "line": 5, + "offset": 4 + }, + "endLine": { + "line": 11, + "offset": 5 + } + }, + "moduleName": "kafka" + }, + "properties": { + "wrapperTypeName": { + "value": "KafkaAnydataConsumer", + "enabled": false, + "editable": false, + "optional": false, + "advanced": false + }, + "payloadFieldName": { + "value": "value", + "enabled": false, + "editable": false, + "optional": false, + "advanced": false + }, + "canDataBind": { + "value": "true", + "enabled": false, + "editable": false, + "optional": false, + "advanced": false + } + } + }, + "output": { + "kafka_databind/main.bal": [ + { + "range": { + "start": { + "line": 5, + "character": 36 + }, + "end": { + "line": 5, + "character": 83 + } + }, + "newText": "(KafkaAnydataConsumer1[] messages) returns error? " + } + ], + "kafka_databind/types.bal": [ + { + "range": { + "start": { + "line": 0, + "character": 0 + }, + "end": { + "line": 0, + "character": 0 + } + }, + "newText": "\nimport ballerina/ai;\n\n" + }, + { + "range": { + "start": { + "line": 5, + "character": 3 + }, + "end": { + "line": 5, + "character": 3 + } + }, + "newText": "\n\ntype KafkaAnydataConsumer1 record {|\n *kafka:AnydataConsumerRecord;\n ai:Schema value;\n|};" + } + ] + } +} diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/source/kafka_databind/Ballerina.toml b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/source/kafka_databind/Ballerina.toml new file mode 100644 index 0000000000..1c194e4030 --- /dev/null +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/test/resources/update_function/source/kafka_databind/Ballerina.toml @@ -0,0 +1,4 @@ +[package] +org = "ballerina" +name = "service_designer" +version = "0.1.0" From 425da379148f0a29cd98296f30d3585ee8faef95 Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Mon, 26 Jan 2026 23:25:59 +0530 Subject: [PATCH 7/8] Preserve the import adding order --- .../servicemodelgenerator/extension/util/DatabindUtil.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java index 7a9ddaf751..f7be94d04c 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java @@ -58,6 +58,7 @@ import java.nio.file.Paths; import java.util.ArrayList; import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -894,7 +895,7 @@ private static String generateNewDataBindTypeName(String contextFilePath, Worksp */ private static Set extractRequiredImports(String baseType, ModulePartNode modulePartNode, Map importsForTypeDef) { - Set imports = new HashSet<>(); + Set imports = new LinkedHashSet<>(); if (baseType.contains(COLON)) { String moduleName = baseType.substring(0, baseType.indexOf(COLON)); From 323f56f6c62301c80bedf14808d4fe45a05bb845 Mon Sep 17 00:00:00 2001 From: LakshanWeerasinghe Date: Tue, 27 Jan 2026 09:20:59 +0530 Subject: [PATCH 8/8] Fix copilot suggestion --- .../extension/util/DatabindUtil.java | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java index f7be94d04c..7fa6948c95 100644 --- a/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java +++ b/service-model-generator/modules/service-model-generator-ls-extension/src/main/java/io/ballerina/servicemodelgenerator/extension/util/DatabindUtil.java @@ -907,14 +907,16 @@ private static Set extractRequiredImports(String baseType, ModulePartNod } } - importsForTypeDef.values().forEach(moduleId -> { - String[] importParts = moduleId.split("/"); - String orgName = importParts[0]; - String moduleName = importParts[1].split(":")[0]; - if (!importExists(modulePartNode, orgName, moduleName)) { - imports.add(getImportStmt(orgName, moduleName)); - } - }); + if (importsForTypeDef != null) { + importsForTypeDef.values().forEach(moduleId -> { + String[] importParts = moduleId.split("/"); + String orgName = importParts[0]; + String moduleName = importParts[1].split(":")[0]; + if (!importExists(modulePartNode, orgName, moduleName)) { + imports.add(getImportStmt(orgName, moduleName)); + } + }); + } return imports; }