diff --git a/flytekit/_ast/__init__.py b/flytekit/_ast/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/flytekit/_ast/parser.py b/flytekit/_ast/parser.py
new file mode 100644
index 0000000000..f2311d811c
--- /dev/null
+++ b/flytekit/_ast/parser.py
@@ -0,0 +1,25 @@
+import ast
+import inspect
+import typing
+
+
+def get_function_param_location(func: typing.Callable, param_name: str) -> (int, int):
+    """
+    Get the line and column number of the parameter in the source code of the function definition.
+    """
+    # Get source code of the function
+    source_lines, start_line = inspect.getsourcelines(func)
+    source_code = "".join(source_lines)
+
+    # Parse the source code into an AST
+    module = ast.parse(source_code)
+
+    # Traverse the AST to find the function definition
+    for node in ast.walk(module):
+        if isinstance(node, ast.FunctionDef) and node.name == func.__name__:
+            for i, arg in enumerate(node.args.args):
+                if arg.arg == param_name:
+                    # Calculate the line and column number of the parameter
+                    line_number = start_line + node.lineno - 1
+                    column_offset = arg.col_offset
+                    return line_number, column_offset
diff --git a/flytekit/clis/sdk_in_container/utils.py b/flytekit/clis/sdk_in_container/utils.py
index 1383e4db6e..5b89870d45 100644
--- a/flytekit/clis/sdk_in_container/utils.py
+++ b/flytekit/clis/sdk_in_container/utils.py
@@ -13,7 +13,8 @@
 
 from flytekit.core.constants import SOURCE_CODE
 from flytekit.exceptions.base import FlyteException
-from flytekit.exceptions.user import FlyteInvalidInputException
+from flytekit.exceptions.user import FlyteCompilationException, FlyteInvalidInputException
+from flytekit.exceptions.utils import annotate_exception_with_code
 from flytekit.loggers import get_level_from_cli_verbosity, logger
 
 project_option = click.Option(
@@ -130,12 +131,14 @@ def pretty_print_traceback(e: Exception, verbosity: int = 1):
     else:
         raise ValueError(f"Verbosity level must be between 0 and 2. Got {verbosity}")
 
-    if hasattr(e, SOURCE_CODE):
-        # TODO: Use other way to check if the background is light or dark
-        theme = "emacs" if "LIGHT_BACKGROUND" in os.environ else "monokai"
-        syntax = Syntax(getattr(e, SOURCE_CODE), "python", theme=theme, background_color="default")
-        panel = Panel(syntax, border_style="red", title=type(e).__name__, title_align="left")
-        console.print(panel, no_wrap=False)
+    if isinstance(e, FlyteCompilationException):
+        e = annotate_exception_with_code(e, e.fn, e.param_name)
+        if hasattr(e, SOURCE_CODE):
+            # TODO: Use other way to check if the background is light or dark
+            theme = "emacs" if "LIGHT_BACKGROUND" in os.environ else "monokai"
+            syntax = Syntax(getattr(e, SOURCE_CODE), "python", theme=theme, background_color="default")
+            panel = Panel(syntax, border_style="red", title=e._ERROR_CODE, title_align="left")
+            console.print(panel, no_wrap=False)
 
 
 def pretty_print_exception(e: Exception, verbosity: int = 1):
@@ -161,20 +164,14 @@ def pretty_print_exception(e: Exception, verbosity: int = 1):
                 pretty_print_grpc_error(cause)
             else:
                 pretty_print_traceback(e, verbosity)
+        else:
+            pretty_print_traceback(e, verbosity)
         return
 
     if isinstance(e, grpc.RpcError):
         pretty_print_grpc_error(e)
         return
 
-    if isinstance(e, AssertionError):
-        click.secho(f"Assertion Error: {e}", fg="red")
-        return
-
-    if isinstance(e, ValueError):
-        click.secho(f"Value Error: {e}", fg="red")
-        return
-
     pretty_print_traceback(e, verbosity)
 
 
diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py
index 25aeb9a1b4..65fd4fed6a 100644
--- a/flytekit/core/interface.py
+++ b/flytekit/core/interface.py
@@ -12,11 +12,16 @@
 
 from flytekit.core import context_manager
 from flytekit.core.artifact import Artifact, ArtifactIDSpecification, ArtifactQuery
+from flytekit.core.context_manager import FlyteContextManager
 from flytekit.core.docstring import Docstring
 from flytekit.core.sentinel import DYNAMIC_INPUT_BINDING
 from flytekit.core.type_engine import TypeEngine, UnionTransformer
-from flytekit.exceptions.user import FlyteValidationException
-from flytekit.exceptions.utils import annotate_exception_with_code
+from flytekit.core.utils import has_return_statement
+from flytekit.exceptions.user import (
+    FlyteMissingReturnValueException,
+    FlyteMissingTypeException,
+    FlyteValidationException,
+)
 from flytekit.loggers import developer_logger, logger
 from flytekit.models import interface as _interface_models
 from flytekit.models.literals import Literal, Scalar, Void
@@ -375,6 +380,18 @@ def transform_function_to_interface(fn: typing.Callable, docstring: Optional[Doc
     signature = inspect.signature(fn)
     return_annotation = type_hints.get("return", None)
 
+    ctx = FlyteContextManager.current_context()
+    # Only check if the task/workflow has a return statement at compile time locally.
+    if (
+        ctx.execution_state
+        and ctx.execution_state.mode is None
+        and return_annotation
+        and type(None) not in get_args(return_annotation)
+        and return_annotation is not type(None)
+        and has_return_statement(fn) is False
+    ):
+        raise FlyteMissingReturnValueException(fn=fn)
+
     outputs = extract_return_annotation(return_annotation)
     for k, v in outputs.items():
         outputs[k] = v  # type: ignore
@@ -382,8 +399,7 @@ def transform_function_to_interface(fn: typing.Callable, docstring: Optional[Doc
     for k, v in signature.parameters.items():  # type: ignore
         annotation = type_hints.get(k, None)
         if annotation is None:
-            err_msg = f"'{k}' has no type. Please add a type annotation to the input parameter."
-            raise annotate_exception_with_code(TypeError(err_msg), fn, k)
+            raise FlyteMissingTypeException(fn=fn, param_name=k)
         default = v.default if v.default is not inspect.Parameter.empty else None
         # Inputs with default values are currently ignored, we may want to look into that in the future
         inputs[k] = (annotation, default)  # type: ignore
diff --git a/flytekit/core/utils.py b/flytekit/core/utils.py
index 4c064e8f34..ca3553e79b 100644
--- a/flytekit/core/utils.py
+++ b/flytekit/core/utils.py
@@ -1,8 +1,10 @@
 import datetime
+import inspect
 import os
 import shutil
 import tempfile
 import time
+import typing
 from abc import ABC, abstractmethod
 from functools import wraps
 from hashlib import sha224 as _sha224
@@ -381,3 +383,13 @@ def get_extra_config(self):
         Get the config of the decorator.
         """
         pass
+
+
+def has_return_statement(func: typing.Callable) -> bool:
+    source_lines = inspect.getsourcelines(func)[0]
+    for line in source_lines:
+        if "return" in line.strip():
+            return True
+        if "yield" in line.strip():
+            return True
+    return False
diff --git a/flytekit/exceptions/user.py b/flytekit/exceptions/user.py
index 1ed0954421..a4b5caa75a 100644
--- a/flytekit/exceptions/user.py
+++ b/flytekit/exceptions/user.py
@@ -97,3 +97,25 @@ def __init__(self, request: typing.Any):
 
 class FlytePromiseAttributeResolveException(FlyteAssertion):
     _ERROR_CODE = "USER:PromiseAttributeResolveError"
+
+
+class FlyteCompilationException(FlyteUserException):
+    _ERROR_CODE = "USER:CompileError"
+
+    def __init__(self, fn: typing.Callable, param_name: typing.Optional[str] = None):
+        self.fn = fn
+        self.param_name = param_name
+
+
+class FlyteMissingTypeException(FlyteCompilationException):
+    _ERROR_CODE = "USER:MissingTypeError"
+
+    def __str__(self):
+        return f"'{self.param_name}' has no type. Please add a type annotation to the input parameter."
+
+
+class FlyteMissingReturnValueException(FlyteCompilationException):
+    _ERROR_CODE = "USER:MissingReturnValueError"
+
+    def __str__(self):
+        return f"{self.fn.__name__} function must return a value. Please add a return statement at the end of the function."
diff --git a/flytekit/exceptions/utils.py b/flytekit/exceptions/utils.py
index cd94ae7002..9b46cb405f 100644
--- a/flytekit/exceptions/utils.py
+++ b/flytekit/exceptions/utils.py
@@ -1,44 +1,28 @@
-import ast
 import inspect
 import typing
 
+from flytekit._ast.parser import get_function_param_location
 from flytekit.core.constants import SOURCE_CODE
+from flytekit.exceptions.user import FlyteUserException
 
 
-def get_function_param_location(func: typing.Callable, param_name: str) -> (int, int):
-    """
-    Get the line and column number of the parameter in the source code of the function definition.
-    """
-    # Get source code of the function
-    source_lines, start_line = inspect.getsourcelines(func)
-    source_code = "".join(source_lines)
-
-    # Parse the source code into an AST
-    module = ast.parse(source_code)
-
-    # Traverse the AST to find the function definition
-    for node in ast.walk(module):
-        if isinstance(node, ast.FunctionDef) and node.name == func.__name__:
-            for i, arg in enumerate(node.args.args):
-                if arg.arg == param_name:
-                    # Calculate the line and column number of the parameter
-                    line_number = start_line + node.lineno - 1
-                    column_offset = arg.col_offset
-                    return line_number, column_offset
-
-
-def get_source_code_from_fn(fn: typing.Callable, param_name: str) -> (str, int):
+def get_source_code_from_fn(fn: typing.Callable, param_name: typing.Optional[str] = None) -> (str, int):
     """
     Get the source code of the function and the column offset of the parameter defined in the input signature.
     """
     lines, start_line = inspect.getsourcelines(fn)
+    if param_name is None:
+        return "".join(f"{start_line + i} {lines[i]}" for i in range(len(lines))), 0
+
     target_line_no, column_offset = get_function_param_location(fn, param_name)
     line_index = target_line_no - start_line
     source_code = "".join(f"{start_line + i} {lines[i]}" for i in range(line_index + 1))
     return source_code, column_offset
 
 
-def annotate_exception_with_code(exception: Exception, fn: typing.Callable, param_name: str) -> Exception:
+def annotate_exception_with_code(
+    exception: FlyteUserException, fn: typing.Callable, param_name: typing.Optional[str] = None
+) -> FlyteUserException:
     """
     Annotate the exception with the source code, and will be printed in the rich panel.
     @param exception: The exception to be annotated.
diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py
index f81031361c..fc57cb7573 100644
--- a/tests/flytekit/integration/remote/test_remote.py
+++ b/tests/flytekit/integration/remote/test_remote.py
@@ -390,14 +390,16 @@ def test_fetch_not_exist_launch_plan(register):
 
 
 def test_execute_reference_task(register):
+    nt = typing.NamedTuple("OutputsBC", [("t1_int_output", int), ("c", str)])
+
     @reference_task(
         project=PROJECT,
         domain=DOMAIN,
         name="basic.basic_workflow.t1",
         version=VERSION,
     )
-    def t1(a: int) -> typing.NamedTuple("OutputsBC", t1_int_output=int, c=str):
-        ...
+    def t1(a: int) -> nt:
+        return nt(t1_int_output=a + 2, c="world")
 
     remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
     execution = remote.execute(
@@ -424,7 +426,7 @@ def test_execute_reference_workflow(register):
         version=VERSION,
     )
     def my_wf(a: int, b: str) -> (int, str):
-        ...
+        return a + 2, b + "world"
 
     remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
     execution = remote.execute(
@@ -451,7 +453,7 @@ def test_execute_reference_launchplan(register):
         version=VERSION,
     )
     def my_wf(a: int, b: str) -> (int, str):
-        ...
+        return 3, "world"
 
     remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
     execution = remote.execute(
diff --git a/tests/flytekit/unit/bin/test_python_entrypoint.py b/tests/flytekit/unit/bin/test_python_entrypoint.py
index 3fd91c2932..079b55ec3b 100644
--- a/tests/flytekit/unit/bin/test_python_entrypoint.py
+++ b/tests/flytekit/unit/bin/test_python_entrypoint.py
@@ -198,6 +198,7 @@ def test_dispatch_execute_user_error_non_recov(mock_write_to_file, mock_upload_d
     def t1(a: int) -> str:
         # Should be interpreted as a non-recoverable user error
         raise ValueError(f"some exception {a}")
+        return "hello"
 
     ctx = context_manager.FlyteContext.current_context()
     with context_manager.FlyteContextManager.with_context(
@@ -242,6 +243,7 @@ def t1(a: int) -> str:
     def my_subwf(a: int) -> typing.List[str]:
         # This also tests the dynamic/compile path
         raise user_exceptions.FlyteRecoverableException(f"recoverable {a}")
+        return ["1", "2"]
 
     ctx = context_manager.FlyteContext.current_context()
     with context_manager.FlyteContextManager.with_context(
diff --git a/tests/flytekit/unit/core/flyte_functools/test_decorators.py b/tests/flytekit/unit/core/flyte_functools/test_decorators.py
index dc4babd8b7..2a2fef233f 100644
--- a/tests/flytekit/unit/core/flyte_functools/test_decorators.py
+++ b/tests/flytekit/unit/core/flyte_functools/test_decorators.py
@@ -76,10 +76,10 @@ def test_unwrapped_task():
     error = completed_process.stderr
     error_str = ""
     for line in error.strip().split("\n"):
-        if line.startswith("TypeError"):
+        if line.startswith("FlyteMissingTypeException"):
             error_str += line
     assert error_str != ""
-    assert error_str.startswith("TypeError: 'args' has no type. Please add a type annotation to the input")
+    assert "'args' has no type. Please add a type annotation" in error_str
 
 
 @pytest.mark.parametrize("script", ["nested_function.py", "nested_wrapped_function.py"])
diff --git a/tests/flytekit/unit/core/test_artifacts.py b/tests/flytekit/unit/core/test_artifacts.py
index c026d1b3ce..2eccdf52d5 100644
--- a/tests/flytekit/unit/core/test_artifacts.py
+++ b/tests/flytekit/unit/core/test_artifacts.py
@@ -332,7 +332,7 @@ def test_basic_option_a3():
 
     @task
     def t3(b_value: str) -> Annotated[pd.DataFrame, a3]:
-        ...
+        return pd.DataFrame({"a": [1, 2, 3], "b": [b_value, b_value, b_value]})
 
     entities = OrderedDict()
     t3_s = get_serializable(entities, serialization_settings, t3)
diff --git a/tests/flytekit/unit/core/test_imperative.py b/tests/flytekit/unit/core/test_imperative.py
index aee88e19d1..f361f748b1 100644
--- a/tests/flytekit/unit/core/test_imperative.py
+++ b/tests/flytekit/unit/core/test_imperative.py
@@ -327,7 +327,7 @@ def ref_t1(
         dataframe: pd.DataFrame,
         imputation_method: str = "median",
     ) -> pd.DataFrame:
-        ...
+        return dataframe
 
     @reference_task(
         project="flytesnacks",
@@ -340,7 +340,7 @@ def ref_t2(
         split_mask: int,
         num_features: int,
     ) -> pd.DataFrame:
-        ...
+        return dataframe
 
     wb = ImperativeWorkflow(name="core.feature_engineering.workflow.fe_wf")
     wb.add_workflow_input("sqlite_archive", FlyteFile[typing.TypeVar("sqlite")])
diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py
index e860729a83..fb0d1e6816 100644
--- a/tests/flytekit/unit/core/test_interface.py
+++ b/tests/flytekit/unit/core/test_interface.py
@@ -162,7 +162,7 @@ def test_parameters_and_defaults():
     ctx = context_manager.FlyteContext.current_context()
 
     def z(a: int, b: str) -> typing.Tuple[int, str]:
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
@@ -172,7 +172,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]:
     assert params.parameters["b"].default is None
 
     def z(a: int, b: str = "hello") -> typing.Tuple[int, str]:
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
@@ -182,7 +182,7 @@ def z(a: int, b: str = "hello") -> typing.Tuple[int, str]:
     assert params.parameters["b"].default.scalar.primitive.string_value == "hello"
 
     def z(a: int = 7, b: str = "eleven") -> typing.Tuple[int, str]:
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
@@ -204,7 +204,7 @@ def z(a: Annotated[int, "some annotation"]) -> Annotated[int, "some annotation"]
     def z(
         a: typing.Optional[int] = None, b: typing.Optional[str] = None, c: typing.Union[typing.List[int], None] = None
     ) -> typing.Tuple[int, str]:
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
@@ -216,7 +216,7 @@ def z(
     assert params.parameters["c"].default.scalar.none_type == Void()
 
     def z(a: int | None = None, b: str | None = None, c: typing.List[int] | None = None) -> typing.Tuple[int, str]:
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
@@ -257,7 +257,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]:
         :param b: bar
         :return: ramen
         """
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z, Docstring(callable_=z))
     typed_interface = transform_interface_to_typed_interface(our_interface)
@@ -282,7 +282,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]:
         out1, out2 : tuple
             ramen
         """
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z, Docstring(callable_=z))
     typed_interface = transform_interface_to_typed_interface(our_interface)
@@ -310,7 +310,7 @@ def z(a: int, b: str) -> typing.NamedTuple("NT", x_str=str, y_int=int):
         y_int : int
             description for y_int
         """
-        ...
+        return 1, "hello world"
 
     our_interface = transform_function_to_interface(z, Docstring(callable_=z))
     typed_interface = transform_interface_to_typed_interface(our_interface)
@@ -338,7 +338,7 @@ def __init__(self, name):
             self.name = name
 
     def z(a: Foo) -> Foo:
-        ...
+        return a
 
     our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
@@ -375,7 +375,7 @@ def t1(a: int) -> int:
 def test_transform_interface_to_list_interface(optional_outputs, expected_type):
     @task
     def t() -> int:
-        ...
+        return 123
 
     list_interface = transform_interface_to_list_interface(t.python_interface, set(), optional_outputs=optional_outputs)
     assert list_interface.outputs["o0"] == typing.List[expected_type]
@@ -395,7 +395,7 @@ def t() -> int:
 def test_map_task_interface(min_success_ratio, expected_type):
     @task
     def t() -> str:
-        ...
+        return "hello"
 
     mt = map_task(t, min_success_ratio=min_success_ratio)
 
diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py
index 2c8f1fa18e..3775c8e12d 100644
--- a/tests/flytekit/unit/core/test_map_task.py
+++ b/tests/flytekit/unit/core/test_map_task.py
@@ -39,7 +39,7 @@ def t2(a: int) -> str:
 
 @task(cache=True, cache_version="1")
 def t3(a: int, b: str, c: float) -> str:
-    pass
+    return "hello"
 
 
 # This test is for documentation.
diff --git a/tests/flytekit/unit/core/test_references.py b/tests/flytekit/unit/core/test_references.py
index 0e6fb9a70d..732b6951d9 100644
--- a/tests/flytekit/unit/core/test_references.py
+++ b/tests/flytekit/unit/core/test_references.py
@@ -50,7 +50,7 @@ def ref_t1(a: typing.List[str]) -> str:
     The interface of the task must match that of the remote task. Otherwise, remote compilation of the workflow will
     fail.
     """
-    ...
+    return "hello"
 
 
 def test_ref():
@@ -81,7 +81,7 @@ def test_ref_task_more():
         version="553018f39e519bdb2597b652639c30ce16b99c79",
     )
     def ref_t1(a: typing.List[str]) -> str:
-        ...
+        return "hello"
 
     @workflow
     def wf1(in1: typing.List[str]) -> str:
@@ -106,7 +106,7 @@ def test_ref_task_more_2():
         version="553018f39e519bdb2597b652639c30ce16b99c79",
     )
     def ref_t1(a: typing.List[str]) -> str:
-        ...
+        return "hello"
 
     @reference_task(
         project="flytesnacks",
@@ -115,7 +115,7 @@ def ref_t1(a: typing.List[str]) -> str:
         version="553018f39e519bdb2597b652639c30ce16b99c79",
     )
     def ref_t2(a: typing.List[str]) -> str:
-        ...
+        return "hello"
 
     @workflow
     def wf1(in1: typing.List[str]) -> str:
@@ -134,6 +134,7 @@ def wf1(in1: typing.List[str]) -> str:
 @reference_workflow(project="proj", domain="development", name="wf_name", version="abc")
 def ref_wf1(a: int) -> typing.Tuple[str, str]:
     ...
+    return "hello", "world"
 
 
 def test_reference_workflow():
@@ -407,7 +408,7 @@ def ref_wf1(p1: str, p2: str) -> None:
 def test_ref_lp_from_decorator():
     @reference_launch_plan(project="project", domain="domain", name="name", version="version")
     def ref_lp1(p1: str, p2: str) -> int:
-        ...
+        return 0
 
     assert ref_lp1.id.name == "name"
     assert ref_lp1.id.project == "project"
@@ -418,9 +419,10 @@ def ref_lp1(p1: str, p2: str) -> int:
 
 
 def test_ref_lp_from_decorator_with_named_outputs():
+    nt = typing.NamedTuple("RefLPOutput", [("o1", int), ("o2", str)])
     @reference_launch_plan(project="project", domain="domain", name="name", version="version")
-    def ref_lp1(p1: str, p2: str) -> typing.NamedTuple("RefLPOutput", o1=int, o2=str):
-        ...
+    def ref_lp1(p1: str, p2: str) -> nt:
+        return nt(o1=1, o2="2")
 
     assert ref_lp1.python_interface.outputs == {"o1": int, "o2": str}
 
@@ -433,7 +435,7 @@ def test_ref_dynamic_task():
         version="553018f39e519bdb2597b652639c30ce16b99c79",
     )
     def ref_t1(a: int) -> str:
-        ...
+        return "hello"
 
     @task
     def t2(a: str, b: str) -> str:
@@ -468,7 +470,7 @@ def test_ref_dynamic_lp():
     def my_subwf(a: int) -> typing.List[int]:
         @reference_launch_plan(project="project", domain="domain", name="name", version="version")
         def ref_lp1(p1: str, p2: str) -> int:
-            ...
+            return 1
 
         s = []
         for i in range(a):
diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py
index 82e4fef245..61988d8501 100644
--- a/tests/flytekit/unit/core/test_serialization.py
+++ b/tests/flytekit/unit/core/test_serialization.py
@@ -12,7 +12,7 @@
 from flytekit.core.python_auto_container import get_registerable_container_image
 from flytekit.core.task import task
 from flytekit.core.workflow import workflow
-from flytekit.exceptions.user import FlyteAssertion
+from flytekit.exceptions.user import FlyteAssertion, FlyteMissingTypeException
 from flytekit.image_spec.image_spec import ImageBuildEngine, _calculate_deduped_hash_from_image_spec
 from flytekit.models.admin.workflow import WorkflowSpec
 from flytekit.models.literals import (
@@ -727,7 +727,7 @@ def wf_with_sub_wf() -> tuple[typing.Optional[int], typing.Optional[int]]:
 
 
 def test_default_args_task_no_type_hint():
-    with pytest.raises(TypeError, match="'a' has no type. Please add a type annotation to the input parameter"):
+    with pytest.raises(FlyteMissingTypeException, match="'a' has no type. Please add a type annotation to the input parameter"):
         @task
         def t1(a=0) -> int:
             return a
diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py
index 8879119eeb..5364a29b88 100644
--- a/tests/flytekit/unit/core/test_type_hints.py
+++ b/tests/flytekit/unit/core/test_type_hints.py
@@ -1351,7 +1351,7 @@ def foo2() -> str:
 
         @task(secret_requests=["test"])
         def foo() -> str:
-            pass
+            return "hello"
 
 
 def test_nested_dynamic():
@@ -1615,6 +1615,7 @@ def run(a: int, b: str) -> typing.Tuple[int, str]:
     @task
     def fail(a: int, b: str) -> typing.Tuple[int, str]:
         raise ValueError("Fail!")
+        return a + 1, b
 
     @task
     def failure_handler(a: int, b: str, err: typing.Optional[FlyteError]) -> typing.Tuple[int, str]:
diff --git a/tests/flytekit/unit/core/test_workflows.py b/tests/flytekit/unit/core/test_workflows.py
index 60daf80af9..43635bcbbb 100644
--- a/tests/flytekit/unit/core/test_workflows.py
+++ b/tests/flytekit/unit/core/test_workflows.py
@@ -14,7 +14,7 @@
 from flytekit.core.condition import conditional
 from flytekit.core.task import task
 from flytekit.core.workflow import WorkflowFailurePolicy, WorkflowMetadata, WorkflowMetadataDefaults, workflow
-from flytekit.exceptions.user import FlyteValidationException, FlyteValueException
+from flytekit.exceptions.user import FlyteValidationException, FlyteValueException, FlyteMissingReturnValueException
 from flytekit.tools.translator import get_serializable
 from flytekit.types.error.error import FlyteError
 
@@ -237,7 +237,7 @@ def no_outputs_wf():
         no_outputs_wf()
 
     # Should raise an exception because it doesn't return something when it should
-    with pytest.raises(AssertionError):
+    with pytest.raises(FlyteMissingReturnValueException):
 
         @workflow
         def one_output_wf() -> int:  # type: ignore