rasa-pro 3.14.0.dev20250818__py3-none-any.whl → 3.14.0.dev20250901__py3-none-any.whl

rasa/core/channels/studio_chat.py

@@ -17,9 +17,9 @@ from typing import (
     Optional,
     Text,
     Tuple,
+    Union,
 )
 
-import orjson
 import structlog
 
 from rasa.core.channels import UserMessage
@@ -47,7 +47,6 @@ if TYPE_CHECKING:
     from sanic import Sanic, Websocket  # type: ignore[attr-defined]
     from socketio import AsyncServer
 
-    from rasa.core.channels.channel import UserMessage
     from rasa.shared.core.trackers import DialogueStateTracker
 
 
@@ -56,7 +55,7 @@ structlogger = structlog.get_logger()
 
 def tracker_as_dump(
     tracker: "DialogueStateTracker", latency: Optional[float] = None
-) -> str:
+) -> Dict[str, Any]:
     """Create a dump of the tracker state."""
     from rasa.shared.core.trackers import get_trackers_for_conversation_sessions
 
@@ -67,11 +66,14 @@ def tracker_as_dump(
     else:
         last_tracker = multiple_tracker_sessions[-1]
 
+    # TODO: this is a bug: the bridge converts this back to json, but it
+    # should be json in the first place
     state = last_tracker.current_state(EventVerbosity.AFTER_RESTART)
 
     if latency is not None:
         state["latency"] = {"rasa_processing_latency_ms": latency}
-    return orjson.dumps(state, option=orjson.OPT_SERIALIZE_NUMPY).decode("utf-8")
+
+    return state
 
 
 def does_need_action_prediction(tracker: "DialogueStateTracker") -> bool:
@@ -153,6 +155,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
         jwt_key: Optional[Text] = None,
         jwt_method: Optional[Text] = "HS256",
         metadata_key: Optional[Text] = "metadata",
+        enable_silence_timeout: bool = False,
     ) -> None:
         """Creates a `StudioChatInput` object."""
         from rasa.core.agent import Agent
@@ -170,6 +173,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
             jwt_key=jwt_key,
             jwt_method=jwt_method,
             metadata_key=metadata_key,
+            enable_silence_timeout=enable_silence_timeout,
         )
 
         # Initialize the Voice Input Channel
@@ -210,14 +214,15 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
             jwt_key=credentials.get("jwt_key"),
             jwt_method=credentials.get("jwt_method", "HS256"),
             metadata_key=credentials.get("metadata_key", "metadata"),
+            enable_silence_timeout=credentials.get("enable_silence_timeout", False),
         )
 
-    async def emit(self, event: str, data: str, room: str) -> None:
+    async def emit(self, event: str, data: Union[Dict, str], room: str) -> None:
         """Emits an event to the websocket."""
-        if not self.sio:
+        if not self.sio_server:
             structlogger.error("studio_chat.emit.sio_not_initialized")
             return
-        await self.sio.emit(event, data, room=room)
+        await self.sio_server.emit(event, data, room=room)
 
     def _register_tracker_update_hook(self) -> None:
         plugin_manager().register(StudioTrackerUpdatePlugin(self))
@@ -250,15 +255,22 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
 
     async def on_message_proxy(
         self,
-        on_new_message: Callable[["UserMessage"], Awaitable[Any]],
-        message: "UserMessage",
+        on_new_message: Callable[[UserMessage], Awaitable[Any]],
+        message: UserMessage,
     ) -> None:
         """Proxies the on_new_message call to the underlying channel.
 
         Triggers a tracker update notification after processing the message.
         """
         self._record_turn_start_time(message.sender_id)
-        await on_new_message(message)
+        try:
+            await on_new_message(message)
+        except Exception as e:
+            structlogger.exception(
+                "studio_chat.on_new_message.error",
+                error=str(e),
+                sender_id=message.sender_id,
+            )
 
         if not self.agent or not self.agent.is_ready():
             structlogger.error("studio_chat.on_message_proxy.agent_not_initialized")
@@ -377,7 +389,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
                 call_state.is_bot_speaking = True
         return ContinueConversationAction()
 
-    def create_output_channel(
+    def _create_output_channel(
         self, voice_websocket: "Websocket", tts_engine: TTSEngine
     ) -> VoiceOutputChannel:
         """Create a voice output channel."""
@@ -407,7 +419,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
 
         # Create a websocket adapter for this connection
         ws_adapter = SocketIOVoiceWebsocketAdapter(
-            sio=self.sio,
+            sio_server=self.sio_server,
             session_id=session_id,
             sid=sid,
             bot_message_evt=self.bot_message_evt,
@@ -455,12 +467,12 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
             task.cancel()
 
     def blueprint(
-        self, on_new_message: Callable[["UserMessage"], Awaitable[Any]]
+        self, on_new_message: Callable[[UserMessage], Awaitable[Any]]
     ) -> SocketBlueprint:
         proxied_on_message = partial(self.on_message_proxy, on_new_message)
         socket_blueprint = super().blueprint(proxied_on_message)
 
-        if not self.sio:
+        if not self.sio_server:
             structlogger.error("studio_chat.blueprint.sio_not_initialized")
             return socket_blueprint
 
@@ -470,12 +482,12 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
         ) -> None:
             self.agent = app.ctx.agent
 
-        @self.sio.on("disconnect", namespace=self.namespace)
+        @self.sio_server.on("disconnect", namespace=self.namespace)
         async def disconnect(sid: Text) -> None:
             structlogger.debug("studio_chat.sio.disconnect", sid=sid)
             self._cleanup_tasks_for_sid(sid)
 
-        @self.sio.on("session_request", namespace=self.namespace)
+        @self.sio_server.on("session_request", namespace=self.namespace)
         async def session_request(sid: Text, data: Optional[Dict]) -> None:
             """Overrides the base SocketIOInput session_request handler.
 
@@ -495,7 +507,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
             if data and data.get("is_voice", False):
                 self._start_voice_session(data["session_id"], sid, proxied_on_message)
 
-        @self.sio.on(self.user_message_evt, namespace=self.namespace)
+        @self.sio_server.on(self.user_message_evt, namespace=self.namespace)
         async def handle_message(sid: Text, data: Dict) -> None:
             """Overrides the base SocketIOInput handle_message handler."""
             # Handle voice messages
@@ -506,10 +518,18 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
                     ws.put_message(data)
                 return
 
-            # Handle text messages
-            await self.handle_user_message(sid, data, proxied_on_message)
+            try:
+                # Handle text messages
+                await self.handle_user_message(sid, data, proxied_on_message)
+            except Exception as e:
+                structlogger.exception(
+                    "studio_chat.sio.handle_message.error",
+                    error=str(e),
+                    sid=sid,
+                )
+                await self.emit("error", str(e), room=sid)
 
-        @self.sio.on("update_tracker", namespace=self.namespace)
+        @self.sio_server.on("update_tracker", namespace=self.namespace)
         async def on_update_tracker(sid: Text, data: Dict) -> None:
             await self.handle_tracker_update(sid, data)
 
@@ -555,9 +575,9 @@ class SocketIOVoiceWebsocketAdapter:
     """Adapter to make Socket.IO work like a Sanic WebSocket for voice channels."""
 
     def __init__(
-        self, sio: "AsyncServer", session_id: str, sid: str, bot_message_evt: str
+        self, sio_server: "AsyncServer", session_id: str, sid: str, bot_message_evt: str
     ) -> None:
-        self.sio = sio
+        self.sio_server = sio_server
         self.bot_message_evt = bot_message_evt
         self._closed = False
         self._receive_queue: asyncio.Queue[Any] = asyncio.Queue()
@@ -576,7 +596,7 @@ class SocketIOVoiceWebsocketAdapter:
     async def send(self, data: Any) -> None:
         """Send data to the client."""
         if not self.closed:
-            await self.sio.emit(self.bot_message_evt, data, room=self.sid)
+            await self.sio_server.emit(self.bot_message_evt, data, room=self.sid)
 
     async def recv(self) -> Any:
         """Receive data from the client."""

rasa/core/channels/voice_stream/voice_channel.py

@@ -11,6 +11,11 @@ from sanic import Websocket  # type: ignore
 from sanic.exceptions import ServerError, WebsocketClosed
 
 from rasa.core.channels import InputChannel, OutputChannel, UserMessage
+from rasa.core.channels.constants import (
+    USER_CONVERSATION_SESSION_END,
+    USER_CONVERSATION_SESSION_START,
+    USER_CONVERSATION_SILENCE_TIMEOUT,
+)
 from rasa.core.channels.voice_ready.utils import (
     CallParameters,
     validate_voice_license_scope,
@@ -48,9 +53,6 @@ from rasa.utils.io import remove_emojis
 logger = structlog.get_logger(__name__)
 
 # define constants for the voice channel
-USER_CONVERSATION_SESSION_END = "/session_end"
-USER_CONVERSATION_SESSION_START = "/session_start"
-USER_CONVERSATION_SILENCE_TIMEOUT = "/silence_timeout"
 
 
 @dataclass

rasa/core/policies/enterprise_search_policy.py

@@ -781,7 +781,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         if not os.path.exists(docs_folder) or not os.path.isdir(docs_folder):
             error_message = (
                 f"Document source directory does not exist or is not a "
-                f"directory: '{docs_folder}'. "
+                f"directory: '{os.path.abspath(docs_folder)}'. "
                 "Please specify a valid path to the documents source directory in the "
                 "vector_store configuration."
             )
@@ -1130,8 +1130,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         embeddings_config: Dict[Text, Any],
         log_source_method: str,
     ) -> None:
-        """
-        Perform the health checks using resolved LLM and embeddings configurations.
+        """Perform the health checks using resolved LLM and embeddings configurations.
         Resolved means the configuration is either:
         - A reference to a model group that has already been expanded into
           its corresponding configuration using the information from
@@ -1160,8 +1159,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
 
     @classmethod
     def get_system_default_prompt_based_on_config(cls, config: Dict[str, Any]) -> str:
-        """
-        Resolves the default prompt template for Enterprise Search Policy based on
+        """Resolves the default prompt template for Enterprise Search Policy based on
         the component's configuration.
 
         - The old prompt is selected when both citation and relevancy check are either
@@ -1192,8 +1190,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         relevancy_check_enabled: bool,
         citation_enabled: bool,
     ) -> str:
-        """
-        Returns the appropriate default prompt template based on the feature flags.
+        """Returns the appropriate default prompt template based on the feature flags.
 
         The selection follows this priority:
         1. If relevancy check is enabled, return the prompt that includes both relevancy

rasa/core/policies/flows/flow_executor.py

@@ -740,7 +740,14 @@ def _run_action_step(
             # do not log about non-existing validation actions of collect steps
             utter_action_name = render_template_variables("{{context.utter}}", context)
             if utter_action_name not in available_actions:
-                structlogger.warning("flow.step.run.action.unknown", action=action_name)
+                structlogger.warning(
+                    "flow.step.run.action.unknown",
+                    action=action_name,
+                    event_info=(
+                        f"The action '{action_name}' is not defined in the domain but "
+                        f"getting triggered by the flow '{step.flow_id}'."
+                    ),
+                )
         return ContinueFlowWithNextStep(events=initial_events)
 
 

rasa/core/run.py

@@ -42,10 +42,20 @@ from rasa.utils import licensing
 logger = logging.getLogger()  # get the root logger
 
 
-def create_http_input_channels(
+def create_input_channels(
     channel: Optional[Text], credentials_file: Optional[Text]
 ) -> List[InputChannel]:
-    """Instantiate the chosen input channel."""
+    """Instantiate the chosen input channel.
+
+    Args:
+        channel (optional): The name of the specific input channel to create.
+        credentials_file: Path to the credentials file containing channel credentials.
+
+    Returns:
+        A list of instantiated input channels. If a specific channel is provided,
+        it returns a list with that single channel. If no channel is specified,
+        it returns a list of all channels defined in the credentials file.
+    """
     if credentials_file:
         all_credentials = read_config_file(credentials_file)
     else:
@@ -253,7 +263,7 @@ def serve_application(
     if not channel and not credentials:
         channel = "cmdline"
 
-    input_channels = create_http_input_channels(channel, credentials)
+    input_channels = create_input_channels(channel, credentials)
 
     if inspect:
         logger.info("Starting development inspector.")

rasa/dialogue_understanding/generator/flow_retrieval.py

@@ -219,11 +219,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
         if self.vector_store is not None:
             with self._model_storage.write_to(self._resource) as model_path:
                 self.vector_store.save_local(model_path)
-        else:
-            structlogger.warning(
-                "flow_retrieval.persist_vector_store.not_initialized",
-                event_info="Vector store is None, not persisted.",
-            )
 
     def _persist_config(self) -> None:
         with self._model_storage.write_to(self._resource) as path:
@@ -249,6 +244,16 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
         )
 
         flows_to_embedd = flows.exclude_link_only_flows()
+
+        if not flows_to_embedd:
+            structlogger.debug(
+                "flow_retrieval.populate_vector_store.no_flows_to_embed",
+                event_info=(
+                    "No flows to embed in the vector store, skipping population."
+                ),
+            )
+            return
+
         embeddings = self._create_embedder(self.config)
         documents = self._generate_flow_documents(flows_to_embedd, domain)
         try:
@@ -420,10 +425,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
             The top k documents with similarity scores.
         """
         if self.vector_store is None:
-            structlogger.error(
-                "flow_retrieval.query_vector_store.vector_store_not_configured",
-                event_info="Vector store is not configured",
-            )
             return []
         try:
             documents_with_scores = (

rasa/dialogue_understanding/patterns/default_flows_for_patterns.yml

@@ -81,7 +81,7 @@ responses:
         rephrase: True
 
   utter_inform_hangup:
-    - text: It seems you are not there anymore. I will hang up shortly.
+    - text: I haven’t heard from you, so I’ll end our conversation shortly.
       metadata:
         rephrase: True
 

rasa/engine/storage/local_model_storage.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
 import logging
+import os
 import shutil
 import sys
+import tarfile
 import tempfile
 import uuid
 from contextlib import contextmanager
 from datetime import datetime
 from pathlib import Path
-from typing import Generator, Optional, Text, Tuple, Union
+from typing import Callable, Generator, Optional, Text, Tuple, Union
 
 from tarsafe import TarSafe
 
@@ -57,6 +59,35 @@ def windows_safe_temporary_directory(
             yield temporary_directory
 
 
+def filter_normpath(member: tarfile.TarInfo, dest_path: str) -> tarfile.TarInfo:
+    """Normalize tar member paths for safe extraction"""
+    if member.name:
+        member.name = os.path.normpath(member.name)
+    return member
+
+
+FilterFunction = Callable[[tarfile.TarInfo, str], Optional[tarfile.TarInfo]]
+
+
+def create_combined_filter(existing_filter: Optional[FilterFunction]) -> FilterFunction:
+    """Create a filter that combines existing filter with path normalization"""
+
+    def combined_filter(
+        member: tarfile.TarInfo, dest_path: str
+    ) -> Optional[tarfile.TarInfo]:
+        """Apply existing filter first, then path normalization"""
+        if existing_filter is not None:
+            filtered_member = existing_filter(member, dest_path)
+            if filtered_member is None:
+                return None  # Rejected by existing filter
+            member = filtered_member  # Use the filtered result
+
+        # Apply our path normalization
+        return filter_normpath(member, dest_path)
+
+    return combined_filter
+
+
 class LocalModelStorage(ModelStorage):
     """Stores and provides output of `GraphComponents` on local disk."""
 
@@ -122,7 +153,19 @@ class LocalModelStorage(ModelStorage):
                 # this restriction in environments where it's not possible
                 # to override this behavior, mostly for internal policy reasons
                 # reference: https://stackoverflow.com/a/49102229
-                tar.extractall(f"\\\\?\\{temporary_directory}")
+                try:
+                    # Use extraction filter to normalize paths for compatibility
+                    # before trying the \\?\ prefix approach first
+                    prev_filter = getattr(tar, "extraction_filter", None)
+                    tar.extraction_filter = create_combined_filter(prev_filter)
+                    tar.extractall(f"\\\\?\\{temporary_directory}")
+                except Exception:
+                    # Fallback for Python versions with tarfile security fix
+                    logger.warning(
+                        "Failed to extract model archive with long path support. "
+                        "Falling back to regular extraction."
+                    )
+                    tar.extractall(temporary_directory)
             else:
                 tar.extractall(temporary_directory)
         LocalModelStorage._assert_not_rasa2_archive(temporary_directory)

rasa/model_manager/model_api.py

@@ -39,11 +39,9 @@ from rasa.model_manager.trainer_service import (
     update_training_status,
 )
 from rasa.model_manager.utils import (
-    InvalidPathException,
     get_logs_content,
     logs_base_path,
     models_base_path,
-    subpath,
 )
 from rasa.model_manager.warm_rasa_process import (
     initialize_warm_rasa_process,
@@ -53,6 +51,7 @@ from rasa.server import ErrorResponse
 from rasa.shared.exceptions import InvalidConfigException
 from rasa.shared.utils.yaml import dump_obj_as_yaml_to_string
 from rasa.studio.upload import build_calm_import_parts
+from rasa.utils.io import InvalidPathException, subpath
 
 dotenv.load_dotenv()
 
@@ -571,10 +570,10 @@ def external_blueprint() -> Blueprint:
     """Create a blueprint for the model manager API."""
     from rasa.core.channels.socketio import SocketBlueprint
 
-    sio = AsyncServer(async_mode="sanic", cors_allowed_origins="*")
-    bp = SocketBlueprint(sio, "", "model_api_external")
+    sio_server = AsyncServer(async_mode="sanic", cors_allowed_origins="*")
+    bp = SocketBlueprint(sio_server, "", "model_api_external")
 
-    create_bridge_server(sio, running_bots)
+    create_bridge_server(sio_server, running_bots)
 
     @bp.get("/health")
     async def health(request: Request) -> response.HTTPResponse:

rasa/model_manager/runner_service.py

@@ -15,11 +15,11 @@ from rasa.model_manager import config
 from rasa.model_manager.utils import (
     logs_path,
     models_base_path,
-    subpath,
     write_encoded_data_to_file,
 )
 from rasa.model_manager.warm_rasa_process import start_rasa_process
 from rasa.studio.prompts import handle_prompts
+from rasa.utils.io import subpath
 
 structlogger = structlog.get_logger()
 

rasa/model_manager/socket_bridge.py

@@ -19,7 +19,7 @@ socket_proxy_clients = {}
 
 
 async def socketio_websocket_traffic_wrapper(
-    sio: AsyncServer,
+    sio_server: AsyncServer,
     running_bots: Dict[str, BotSession],
     sid: str,
     auth: Optional[Dict],
@@ -55,7 +55,9 @@ async def socketio_websocket_traffic_wrapper(
         structlogger.error("model_runner.bot_not_alive", deployment_id=deployment_id)
         raise ConnectionRefusedError("model_runner.bot_not_alive")
 
-    client = await create_bridge_client(sio, bot.internal_url, sid, deployment_id)
+    client = await create_bridge_client(
+        sio_server, bot.internal_url, sid, deployment_id
+    )
 
     if client.sid is not None:
         structlogger.debug(
@@ -70,20 +72,24 @@ async def socketio_websocket_traffic_wrapper(
         raise ConnectionRefusedError("model_runner.bot_connection_failed")
 
 
-def create_bridge_server(sio: AsyncServer, running_bots: Dict[str, BotSession]) -> None:
+def create_bridge_server(
+    sio_server: AsyncServer, running_bots: Dict[str, BotSession]
+) -> None:
     """Create handlers for the socket server side.
 
     Forwards messages coming from the user to the bot.
     """
 
-    @sio.on("connect")
+    @sio_server.on("connect")
     async def socketio_websocket_traffic(
         sid: str, environ: Dict, auth: Optional[Dict]
     ) -> bool:
         """Bridge websockets between user chat socket and bot server."""
-        return await socketio_websocket_traffic_wrapper(sio, running_bots, sid, auth)
+        return await socketio_websocket_traffic_wrapper(
+            sio_server, running_bots, sid, auth
+        )
 
-    @sio.on("disconnect")
+    @sio_server.on("disconnect")
     async def disconnect(sid: str) -> None:
         """Disconnect the bot connection."""
         structlogger.debug("model_runner.bot_disconnect", sid=sid)
@@ -91,7 +97,7 @@ def create_bridge_server(sio: AsyncServer, running_bots: Dict[str, BotSession])
             await socket_proxy_clients[sid].disconnect()
             del socket_proxy_clients[sid]
 
-    @sio.on("*")
+    @sio_server.on("*")
     async def handle_message(event: str, sid: str, data: Dict[str, Any]) -> None:
         """Bridge messages between user and bot.
 
@@ -108,7 +114,7 @@ def create_bridge_server(sio: AsyncServer, running_bots: Dict[str, BotSession])
 
 
 async def create_bridge_client(
-    sio: AsyncServer, url: str, sid: str, deployment_id: str
+    sio_server: AsyncServer, url: str, sid: str, deployment_id: str
 ) -> AsyncClient:
     """Create a new socket bridge client.
 
@@ -123,36 +129,36 @@ async def create_bridge_client(
         structlogger.debug(
             "model_runner.bot_session_confirmed", deployment_id=deployment_id
         )
-        await sio.emit("session_confirm", room=sid)
+        await sio_server.emit("session_confirm", room=sid)
 
     @client.event  # type: ignore[misc]
     async def bot_message(data: Dict[str, Any]) -> None:
         structlogger.debug("model_runner.bot_message", deployment_id=deployment_id)
-        await sio.emit("bot_message", data, room=sid)
+        await sio_server.emit("bot_message", data, room=sid)
 
     @client.event  # type: ignore[misc]
     async def error(data: Dict[str, Any]) -> None:
         structlogger.debug(
             "model_runner.bot_error", deployment_id=deployment_id, data=data
         )
-        await sio.emit("error", data, room=sid)
+        await sio_server.emit("error", data, room=sid)
 
     @client.event  # type: ignore[misc]
     async def tracker(data: Dict[str, Any]) -> None:
-        await sio.emit("tracker", json.loads(data), room=sid)
+        await sio_server.emit("tracker", json.loads(data), room=sid)
 
     @client.event  # type: ignore[misc]
     async def disconnect() -> None:
         structlogger.debug(
             "model_runner.bot_connection_closed", deployment_id=deployment_id
         )
-        await sio.emit("disconnect", room=sid)
+        await sio_server.emit("disconnect", room=sid)
 
     @client.event  # type: ignore[misc]
     async def connect_error() -> None:
         structlogger.error(
             "model_runner.bot_connection_error", deployment_id=deployment_id
         )
-        await sio.emit("disconnect", room=sid)
+        await sio_server.emit("disconnect", room=sid)
 
     return client

rasa/model_manager/trainer_service.py

@@ -14,7 +14,6 @@ from rasa.model_manager.utils import (
     ensure_base_directory_exists,
     logs_path,
     models_base_path,
-    subpath,
     write_encoded_data_to_file,
 )
 from rasa.model_manager.warm_rasa_process import (
@@ -22,6 +21,7 @@ from rasa.model_manager.warm_rasa_process import (
 )
 from rasa.model_training import generate_random_model_name
 from rasa.studio.prompts import handle_prompts
+from rasa.utils.io import subpath
 
 structlogger = structlog.get_logger()
 
@@ -53,6 +53,15 @@ class TrainingSession(BaseModel):
         """Check if the training is running."""
         return self.status == TrainingSessionStatus.RUNNING
 
+    def has_just_finished(self) -> bool:
+        if not self.is_status_indicating_alive():
+            # skip if the training is not running
+            return False
+        if self.process.poll() is None:
+            # process is still running
+            return False
+        return True
+
     def model_path(self) -> str:
         """Return the path to the model."""
         return subpath(models_base_path(), f"{self.model_name}.tar.gz")
@@ -89,14 +98,8 @@ def terminate_training(training: TrainingSession) -> None:
 
 
 def update_training_status(training: TrainingSession) -> None:
-    if not training.is_status_indicating_alive():
-        # skip if the training is not running
-        return
-    if training.process.poll() is None:
-        # process is still running
-        return
-
-    complete_training(training)
+    if training.has_just_finished():
+        complete_training(training)
 
 
 def complete_training(training: TrainingSession) -> None: