Ensure proper cancellation of in-flight workflow tasks (#1819)

Fixes a bug where an in-flight workflow task was not cancelled when the user clicked "Stop Generating" and submitted a new prompt. This left the handler stuck, most visibly when an OAuth login window was blocked by popup blocker or closed by the user.

Closes #1818 

## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing/index.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.



## Summary by CodeRabbit

* **Bug Fixes**
  * Improved workflow cancellation: new runs preempt and await prior runs, reliably clear running-task state, and skip sending completion/observability messages when canceled.
  * Guarded completion handling to avoid removing active conversation handlers for replaced tasks.
  * Refined async context-manager behavior to avoid masking cancellations; raises only for unexpected unfinished runs when no external exception is present.

* **Tests**
  * Added tests for context-exit behavior, cancellation propagation, workflow replacement handoff, and WebSocket workflow lifecycle.

Authors:
  - Patrick Chin (https://github.com/thepatrickchin)

Approvers:
  - Will Killian (https://github.com/willkill07)

URL: #1819

Author: thepatrickchin

packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py

@@ -272,22 +272,33 @@ async def process_workflow_request(self, user_message_as_validated_type: WebSock
             self._initialize_workflow_request(user_message_as_validated_type)
             message_content: typing.Any = await self._process_websocket_user_message(user_message_as_validated_type)
 
-            if (self._running_workflow_task is None):
+            if self._workflow_schema_type is None:
+                raise RuntimeError("Workflow schema type is not initialized")
 
-                def _done_callback(_task: asyncio.Task):
+            if self._running_workflow_task is not None:
+                self._running_workflow_task.cancel()
+                try:
+                    await self._running_workflow_task
+                except (asyncio.CancelledError, Exception):
+                    pass
+                self._running_workflow_task = None
+
+            _conversation_id = self._conversation_id
+
+            def _done_callback(_task: asyncio.Task):
+                if self._running_workflow_task is _task:
                     self._running_workflow_task = None
-                    if self._conversation_id:
-                        self._worker.remove_conversation_handler(self._conversation_id)
-
-                if self._workflow_schema_type is None:
-                    raise RuntimeError("Workflow schema type is not initialized")
-                self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(payload=message_content,
-                                       user_message_id=self._message_parent_id,
-                                       conversation_id=self._conversation_id,
-                                       result_type=self._schema_output_mapping[self._workflow_schema_type],
-                                       output_type=self._schema_output_mapping[self._workflow_schema_type]))
-                self._running_workflow_task.add_done_callback(_done_callback)
+                if self._running_workflow_task is None and _conversation_id and \
+                   self._worker.get_conversation_handler(_conversation_id) is self:
+                    self._worker.remove_conversation_handler(_conversation_id)
+
+            self._running_workflow_task = asyncio.create_task(
+                self._run_workflow(payload=message_content,
+                                   user_message_id=self._message_parent_id,
+                                   conversation_id=self._conversation_id,
+                                   result_type=self._schema_output_mapping[self._workflow_schema_type],
+                                   output_type=self._schema_output_mapping[self._workflow_schema_type]))
+            self._running_workflow_task.add_done_callback(_done_callback)
 
         except ValueError as e:
             logger.exception("User message content not found: %s", str(e))
@@ -439,6 +450,7 @@ async def _run_workflow(self,
                             result_type: type | None = None,
                             output_type: type | None = None) -> None:
 
+        _cancelled = False
         try:
             auth_callback = self._flow_handler.authenticate if self._flow_handler else None
             async with self._session_manager.session(user_id=self._user_id,
@@ -466,6 +478,10 @@ async def _run_workflow(self,
 
                     await self.create_websocket_message(data_model=value, status=WebSocketMessageStatus.IN_PROGRESS)
 
+        except asyncio.CancelledError:
+            _cancelled = True
+            raise
+
         except Exception as e:
             logger.exception("Unhandled workflow error")
             await self.create_websocket_message(data_model=Error(code=ErrorTypes.WORKFLOW_ERROR,
@@ -475,12 +491,16 @@ async def _run_workflow(self,
                                                 status=WebSocketMessageStatus.IN_PROGRESS)
 
         finally:
-            await self.create_websocket_message(data_model=SystemResponseContent(),
-                                                message_type=WebSocketMessageType.RESPONSE_MESSAGE,
-                                                status=WebSocketMessageStatus.COMPLETE)
-
-            # Send observability trace after completion message
-            if self._pending_observability_trace is not None:
-                await self.create_websocket_message(data_model=self._pending_observability_trace,
-                                                    message_type=WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE)
+            try:
+                if not _cancelled:
+                    await self.create_websocket_message(data_model=SystemResponseContent(),
+                                                        message_type=WebSocketMessageType.RESPONSE_MESSAGE,
+                                                        status=WebSocketMessageStatus.COMPLETE)
+
+                    # Send observability trace after completion message
+                    if self._pending_observability_trace is not None:
+                        await self.create_websocket_message(
+                            data_model=self._pending_observability_trace,
+                            message_type=WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE)
+            finally:
                 self._pending_observability_trace = None

packages/nvidia_nat_core/src/nat/runtime/runner.py

@@ -150,7 +150,7 @@ async def __aexit__(self, exc_type, exc_value, traceback):
 
         self._context_state.runtime_type.reset(self._runtime_type_token)
 
-        if (self._state not in (RunnerState.COMPLETED, RunnerState.FAILED)):
+        if (self._state not in (RunnerState.COMPLETED, RunnerState.FAILED)) and exc_type is None:
             raise ValueError("Cannot exit the context without completing the workflow")
 
     @typing.overload

packages/nvidia_nat_core/tests/nat/runtime/test_runner.py

@@ -282,3 +282,78 @@ async def test_runner_state_management():
         async with runner:
             result = await runner.result()
             assert result == "test!"
+
+
+async def test_runner_aexit_raises_on_incomplete_clean_exit():
+    """Test that Runner raises ValueError when exited cleanly without completing the workflow."""
+
+    async with WorkflowBuilder() as builder:
+        entry_fn = await builder.add_function(name="test_function", config=SingleOutputConfig())
+
+        context_state = ContextState()
+        exporter_manager = ExporterManager()
+
+        with pytest.raises(ValueError, match="Cannot exit the context without completing the workflow"):
+            async with Runner(input_message="test",
+                              entry_fn=entry_fn,
+                              context_state=context_state,
+                              exporter_manager=exporter_manager):
+                pass  # exit without calling result()
+
+
+async def test_runner_aexit_allows_cancelled_error_to_propagate():
+    """Test that Runner does not mask CancelledError with a ValueError on exit."""
+    import asyncio
+
+    async with WorkflowBuilder() as builder:
+        entry_fn = await builder.add_function(name="test_function", config=SingleOutputConfig())
+
+        context_state = ContextState()
+        exporter_manager = ExporterManager()
+
+        with pytest.raises(asyncio.CancelledError):
+            async with Runner(input_message="test",
+                              entry_fn=entry_fn,
+                              context_state=context_state,
+                              exporter_manager=exporter_manager):
+                raise asyncio.CancelledError()
+
+
+async def test_runner_workflow_replacement_handoff():
+    """Test the workflow-replacement handoff path tied to the message_handler regression.
+
+    Cancelling the first in-flight Runner via task.cancel() must not mask CancelledError
+    with a ValueError (runner.py fix), and the immediately-following second Runner must
+    run to completion on the same context_state/exporter_manager.
+    """
+    import asyncio
+
+    async with WorkflowBuilder() as builder:
+        entry_fn = await builder.add_function(name="test_function", config=SingleOutputConfig())
+
+        context_state = ContextState()
+        exporter_manager = ExporterManager()
+
+        # Simulate message_handler cancelling an in-flight task when a new message arrives.
+        async def _first_workflow():
+            async with Runner(input_message="first",
+                              entry_fn=entry_fn,
+                              context_state=context_state,
+                              exporter_manager=exporter_manager):
+                await asyncio.sleep(0)  # yield so external cancel can be delivered
+
+        first_task = asyncio.create_task(_first_workflow())
+        await asyncio.sleep(0)  # let the task enter the Runner context
+        first_task.cancel()
+
+        # CancelledError must propagate cleanly — not be masked by ValueError.
+        with pytest.raises(asyncio.CancelledError):
+            await first_task
+
+        # Handoff: second Runner starts immediately on the same context and runs to completion.
+        async with Runner(input_message="second",
+                          entry_fn=entry_fn,
+                          context_state=context_state,
+                          exporter_manager=exporter_manager) as runner2:
+            result = await runner2.result()
+            assert result == "second!"

packages/nvidia_nat_core/tests/nat/server/test_unified_api_server.py

@@ -18,6 +18,7 @@
 import json
 import os
 import re
+from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock
 from unittest.mock import MagicMock
 from unittest.mock import patch
@@ -43,6 +44,7 @@
 from nat.data_models.api_server import ErrorTypes
 from nat.data_models.api_server import ObservabilityTraceContent
 from nat.data_models.api_server import ResponseIntermediateStep
+from nat.data_models.api_server import ResponseObservabilityTrace
 from nat.data_models.api_server import ResponsePayloadOutput
 from nat.data_models.api_server import SystemIntermediateStepContent
 from nat.data_models.api_server import SystemResponseContent
@@ -1053,3 +1055,130 @@ async def test_restore_execution_state_sends_prompt_with_remaining_timeout():
     call_kwargs = handler.create_websocket_message.call_args[1]
     sent_content = call_kwargs["data_model"]
     assert sent_content.timeout == 7
+
+
+async def test_process_workflow_request_cancels_in_flight_task():
+    """A new workflow request cancels any in-flight task before creating a replacement."""
+    mock_socket = AsyncMock()
+    mock_session_manager = MagicMock()
+    mock_step_adaptor = MagicMock()
+    mock_worker = MagicMock()
+    mock_worker.get_conversation_handler.return_value = None
+
+    handler = WebSocketMessageHandler(
+        socket=mock_socket,
+        session_manager=mock_session_manager,
+        step_adaptor=mock_step_adaptor,
+        worker=mock_worker,
+    )
+
+    existing_task = asyncio.create_task(asyncio.sleep(100))
+    handler._running_workflow_task = existing_task
+
+    msg = WebSocketUserMessage.model_validate({**user_message, "type": "user_message"})
+
+    async def _noop_workflow(*args, **kwargs):
+        await asyncio.sleep(0)
+
+    with patch.object(handler, "_run_workflow", _noop_workflow):
+        await handler.process_workflow_request(msg)
+
+    assert existing_task.cancelled()
+    assert handler._running_workflow_task is not None
+    assert handler._running_workflow_task is not existing_task
+
+    new_task = handler._running_workflow_task
+    new_task.cancel()
+    try:
+        await new_task
+    except (asyncio.CancelledError, Exception):
+        pass
+
+
+async def test_done_callback_guards_against_stale_task():
+    """_done_callback does not clear _running_workflow_task when the task has been replaced."""
+    mock_socket = AsyncMock()
+    mock_session_manager = MagicMock()
+    mock_step_adaptor = MagicMock()
+    mock_worker = MagicMock()
+    mock_worker.get_conversation_handler.return_value = None
+
+    handler = WebSocketMessageHandler(
+        socket=mock_socket,
+        session_manager=mock_session_manager,
+        step_adaptor=mock_step_adaptor,
+        worker=mock_worker,
+    )
+
+    msg = WebSocketUserMessage.model_validate({**user_message, "type": "user_message"})
+    completed = asyncio.Event()
+
+    async def _quick_workflow(*args, **kwargs):
+        await asyncio.sleep(0)
+        completed.set()
+
+    with patch.object(handler, "_run_workflow", _quick_workflow):
+        await handler.process_workflow_request(msg)
+
+    # Simulate a second request replacing the first task reference
+    second_task = asyncio.create_task(asyncio.sleep(100))
+    handler._running_workflow_task = second_task
+
+    # Let the first task complete and fire its done callback
+    await completed.wait()
+    await asyncio.sleep(0)
+
+    # second_task must remain untouched by the first task's callback
+    assert handler._running_workflow_task is second_task
+    mock_worker.remove_conversation_handler.assert_not_called()
+
+    second_task.cancel()
+    try:
+        await second_task
+    except asyncio.CancelledError:
+        pass
+
+
+async def test_run_workflow_skips_response_on_cancellation():
+    """When _run_workflow is cancelled, RESPONSE_MESSAGE is not sent and pending trace is cleared."""
+    mock_socket = AsyncMock()
+    mock_session_manager = MagicMock()
+    mock_step_adaptor = MagicMock()
+    mock_worker = MagicMock()
+
+    handler = WebSocketMessageHandler(
+        socket=mock_socket,
+        session_manager=mock_session_manager,
+        step_adaptor=mock_step_adaptor,
+        worker=mock_worker,
+    )
+    handler.create_websocket_message = AsyncMock()
+    handler._user_message_payload = {}
+
+    handler._pending_observability_trace = ResponseObservabilityTrace(observability_trace_id="trace-to-clear")
+
+    blocking_event = asyncio.Event()
+
+    @asynccontextmanager
+    async def _mock_session(*args, **kwargs):
+        yield MagicMock()
+
+    mock_session_manager.session = _mock_session
+
+    async def _blocking_generator(*args, **kwargs):
+        blocking_event.set()
+        await asyncio.sleep(100)
+        yield  # pragma: no cover
+
+    with patch("nat.front_ends.fastapi.message_handler.generate_streaming_response", _blocking_generator):
+        task = asyncio.create_task(handler._run_workflow(payload="test"))
+        await blocking_event.wait()
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+    for call in handler.create_websocket_message.call_args_list:
+        msg_type = call.kwargs.get("message_type") or (call.args[1] if len(call.args) > 1 else None)
+        assert msg_type != WebSocketMessageType.RESPONSE_MESSAGE
+
+    assert handler._pending_observability_trace is None