diff --git a/ets2panda/checker/ETSchecker.h b/ets2panda/checker/ETSchecker.h
index c512dfb6581254cad68adecaa950753e682689b4..092499702c63365c9b3ec8bd69239565a00cdccf 100644
--- a/ets2panda/checker/ETSchecker.h
+++ b/ets2panda/checker/ETSchecker.h
@@ -297,6 +297,7 @@ public:
     bool NeedTypeInference(const ir::ScriptFunction *lambda);
     std::vector<bool> FindTypeInferenceArguments(const ArenaVector<ir::Expression *> &arguments);
     void InferTypesForLambda(ir::ScriptFunction *lambda, ir::ETSFunctionType *calleeType);
+    void HandleLambdaTypeInfer(ir::AstNode *typeAnn, ir::ScriptFunction *const lambda);
     bool TypeInference(Signature *signature, const ArenaVector<ir::Expression *> &arguments,
                        TypeRelationFlag flags = TypeRelationFlag::NONE);
     bool CheckLambdaAssignable(ir::Expression *param, ir::ScriptFunction *lambda);
diff --git a/ets2panda/checker/ets/typeCheckingHelpers.cpp b/ets2panda/checker/ets/typeCheckingHelpers.cpp
index ca5b3f41ca931c838456a7bbeca39e637dbbae55..52d8b95f6ae630bfdb345e2a0c4464dfac3eb0c8 100644
--- a/ets2panda/checker/ets/typeCheckingHelpers.cpp
+++ b/ets2panda/checker/ets/typeCheckingHelpers.cpp
@@ -931,8 +931,8 @@ bool ETSChecker::TypeInference(Signature *signature, const ArenaVector<ir::Expre
             typeAnn = DerefETSTypeReference(typeAnn);
         }
 
-        ASSERT(typeAnn->IsETSFunctionType());
-        InferTypesForLambda(lambda, typeAnn->AsETSFunctionType());
+        HandleLambdaTypeInfer(typeAnn, lambda);
+
         Type *const argumentType = GetApparentType(arrowFuncExpr->Check(this));
         Type *const parameterType = GetApparentType(signature->Params()[index]->TsType());
         const Type *targetType = TryGettingFunctionTypeFromInvokeFunction(parameterType);
@@ -947,6 +947,23 @@ bool ETSChecker::TypeInference(Signature *signature, const ArenaVector<ir::Expre
     return invocable;
 }
 
+void ETSChecker::HandleLambdaTypeInfer(ir::AstNode *typeAnn, ir::ScriptFunction *const lambda)
+{
+    ASSERT(typeAnn->IsETSFunctionType() || typeAnn->IsETSUnionType());
+
+    if (typeAnn->IsETSFunctionType()) {
+        InferTypesForLambda(lambda, typeAnn->AsETSFunctionType());
+        return;
+    }
+
+    for (auto type : typeAnn->AsETSUnionType()->Types()) {
+        if (type->IsETSFunctionType()) {
+            InferTypesForLambda(lambda, type->AsETSFunctionType());
+            return;
+        }
+    }
+}
+
 bool ETSChecker::ExtensionETSFunctionType(checker::Type *type)
 {
     if (!type->IsETSFunctionType()) {
diff --git a/ets2panda/test/runtime/ets/optional_trailing_lambda.ets b/ets2panda/test/runtime/ets/optional_trailing_lambda.ets
new file mode 100644
index 0000000000000000000000000000000000000000..361f40cd87abec090a9c1464e831ea25faf4ae9d
--- /dev/null
+++ b/ets2panda/test/runtime/ets/optional_trailing_lambda.ets
@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2024 Huawei Device Co., Ltd.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+let a = 0;
+
+function foo(exec?: () => void): void {
+    exec?.();
+    a += 1;
+}
+
+function bar() {
+    foo() {
+        a += 1;
+        assert a == 1;
+        foo() {
+            a += 1;
+            assert a == 2;
+            foo();
+        }
+    }
+}
+
+export function main() {
+    a = 0
+    bar();
+    assert a == 5;
+}