From ea17a4d3c368d4aff4db3631ec53f78805955cfc Mon Sep 17 00:00:00 2001
From: Vincent <0426vincent@gmail.com>
Date: Sun, 5 Jan 2025 20:48:54 -0800
Subject: [PATCH] fix: add FlyteValueException handling and clean up
 exit_handler calls

Signed-off-by: Vincent <0426vincent@gmail.com>
---
 flytekit/core/base_task.py                         | 14 +++++++++++++-
 flytekit/interactive/vscode_lib/decorator.py       |  2 +-
 .../flyteinteractive/jupyter_lib/decorator.py      |  2 +-
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py
index 6430aa9eac..2a314e1aab 100644
--- a/flytekit/core/base_task.py
+++ b/flytekit/core/base_task.py
@@ -75,7 +75,10 @@
     FlyteNonRecoverableSystemException,
     FlyteUploadDataException,
 )
-from flytekit.exceptions.user import FlyteUserRuntimeException
+from flytekit.exceptions.user import (
+    FlyteUserRuntimeException,
+    FlyteValueException,
+)
 from flytekit.loggers import logger
 from flytekit.models import dynamic_job as _dynamic_job
 from flytekit.models import interface as _interface_models
@@ -770,6 +773,15 @@ def dispatch_execute(
             ):
                 return native_outputs
 
+            if isinstance(native_outputs, VoidPromise):
+                return _literal_models.LiteralMap(literals={})
+
+            if native_outputs is not None and len(list(self._outputs_interface.keys())) == 0:
+                raise FlyteValueException(
+                    native_outputs,
+                    f"Interface has {len(self.python_interface.outputs)} outputs.",
+                )
+
             try:
                 with timeit("dispatch execute"):
                     literals_map, native_outputs_as_map = run_sync(
diff --git a/flytekit/interactive/vscode_lib/decorator.py b/flytekit/interactive/vscode_lib/decorator.py
index 055bda6639..bfa8277531 100644
--- a/flytekit/interactive/vscode_lib/decorator.py
+++ b/flytekit/interactive/vscode_lib/decorator.py
@@ -456,7 +456,7 @@ def execute(self, *args, **kwargs):
         # 5. Prepare the launch.json
         prepare_launch_json()
 
-        return exit_handler(
+        exit_handler(
             child_process=child_process,
             task_function=self.task_function,
             args=args,
diff --git a/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/jupyter_lib/decorator.py b/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/jupyter_lib/decorator.py
index d4ba5ee165..34c474d607 100644
--- a/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/jupyter_lib/decorator.py
+++ b/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/jupyter_lib/decorator.py
@@ -168,7 +168,7 @@ def execute(self, *args, **kwargs):
 
         write_example_notebook(task_function=self.task_function, notebook_dir=self.notebook_dir)
 
-        return exit_handler(
+        exit_handler(
             child_process=child_process,
             task_function=self.task_function,
             args=args,