QTZPOC-20

QTZPOC-20: Refactor ModelImport status engine for lazy initialization, add byte-counted size validation to bundle extraction, optimize MLflow run ID lookups, and ensure workflow cancellation propagates correctly.

laborious/activities/model_import.py

@@ -40,7 +40,7 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.repository.mongodb_repository import MongoDBRepository
     from sientia_do.temporal.constants import now
     from sqlalchemy import create_engine, text
-    from sqlalchemy.engine import URL
+    from sqlalchemy.engine import URL, Engine
     from sqlalchemy.exc import DataError, IntegrityError, ProgrammingError
     from sqlalchemy.sql import Executable
     from temporalio.common import RetryPolicy
@@ -321,25 +321,37 @@ class ModelImport(SientiaMonitoring):
         self.import_minio_repository = minio_repository
         self.import_mlflow_repository = mlflow_repository
         self.mongo_repository = mongo_repository
-        self.status_engine = create_engine(
-            URL.create(
-                'postgresql',
-                username=status_config['user'],
-                password=status_config['password'],
-                host=status_config['host'],
-                port=status_config['port'],
-                database=status_config['dbname'],
-            ),
-            pool_size=2,
-            max_overflow=3,
-            pool_pre_ping=True,
-        )
+        self._status_engine: Engine | None = None
+
+    @property
+    def status_engine(self) -> Engine:
+        """The BFF status database's engine, created on first access and cached from then on.
+
+        A worker pod where the import feature is never exercised never opens a connection pool
+        against that Postgres: nothing in this module reaches `status_engine` unless an import
+        activity actually runs.
+        """
+        if self._status_engine is None:
+            self._status_engine = create_engine(
+                URL.create(
+                    'postgresql',
+                    username=self.status_config['user'],
+                    password=self.status_config['password'],
+                    host=self.status_config['host'],
+                    port=self.status_config['port'],
+                    database=self.status_config['dbname'],
+                ),
+                pool_size=2,
+                max_overflow=3,
+                pool_pre_ping=True,
+            )
+        return self._status_engine
 
     def close(self) -> None:
-        """Dispose the status engine. The import holds no other connection."""
-        engine = getattr(self, 'status_engine', None)
-        if engine is not None:
-            engine.dispose()
+        """Dispose the status engine, when one was ever created. The import holds no other
+        connection."""
+        if self._status_engine is not None:
+            self._status_engine.dispose()
 
     # ------------------------------------------------------------------ helpers
 
@@ -1128,7 +1140,7 @@ class ModelImport(SientiaMonitoring):
             upsert=True,
             return_document=ReturnDocument.AFTER,
         )
-        return str(doc['seq']), 0
+        return str(doc['seq']), 0  # type: ignore[index]
 
     # ----------------------------------------------------------- cleanup and retention
 
@@ -1217,7 +1229,12 @@ class ModelImport(SientiaMonitoring):
             # The read still runs every time (it is how drift — someone editing the rule outside
             # this code — gets detected); only the write is skipped when nothing has drifted.
             self.info(f'import retention rule already correct on {bucket} ({days} days)', metadata)
-            return {'bucket': bucket, 'prefix': prefix, 'days': days, 'rules': len(existing.rules)}
+            return {
+                'bucket': bucket,
+                'prefix': prefix,
+                'days': days,
+                'rules': len(existing.rules),  # type: ignore[union-attr]
+            }
 
         rules = [rule for rule in (existing.rules if existing else []) if rule.rule_id != rule_id]
         rules.append(

laborious/activities/opc.py

@@ -106,13 +106,21 @@ class OPC(SientiaMonitoring):
         4. Logs connection status for operational visibility
 
         Raises:
-            Exception: If OPC repository initialization fails or connection
-                establishment encounters critical errors
+            ValueError: A server whose `cert_path` is set but whose `server_uri`
+                (or private key) is missing fails `OpcRepository.set_security()`
+                with no try/except between here and `main()`. This is a
+                configuration error, not a connectivity one, and is intentionally
+                fatal to the whole worker process — fix the deployment config
+                rather than catch it here.
 
         Note:
-            Connection failures are logged and reported but do not prevent
-            the initialization of other OPC servers. Each server is handled
-            independently to ensure maximum availability.
+            Connection failures (server unreachable, timeout, etc.) are logged
+            and reported but do not prevent the initialization of other OPC
+            servers — each server's `connect()` call catches those on its own.
+            This isolation does NOT extend to configuration errors raised before
+            a connection is attempted (see `Raises` above): a single
+            misconfigured server intentionally stops every other server, and the
+            whole worker, from starting.
         """
         self.logger.info('Initializing OPC servers...')
         for opc_id, server in self.opc_servers.items():

laborious/utils/bundle/reader.py

@@ -35,6 +35,9 @@ import nacl.exceptions
 from laborious.utils.bundle import format as fmt
 from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep
 
+# Gate 5 — chunk size for the byte-counted extraction copy loop.
+_EXTRACT_CHUNK_BYTES = 1024 * 1024
+
 # Gate 6 — the layout `bundle.build_zip` writes and `model_repository.py` reads back.
 METADATA_FILE_NAME = 'metadata.json'
 ARTIFACTS_DIR_NAME = 'artifacts'
@@ -389,8 +392,23 @@ class BundleReader:
         resolved.parent.mkdir(parents=True, exist_ok=True)
         # Written as a plain file: the archive's stored permissions are never restored, so no
         # member can arrive executable.
+        #
+        # Copied in chunks, counting bytes actually written, instead of `shutil.copyfileobj`:
+        # `entry.file_size` is central-directory metadata gate 4 already checked in aggregate, but a
+        # DEFLATE stream can be crafted to decompress to far more bytes than its own header
+        # declares. Aborting as soon as the running count exceeds `file_size` catches that member
+        # before the excess ever reaches disk, rather than trusting the declared size alone.
+        written = 0
         with archive.open(entry) as source, resolved.open('wb') as sink:
-            shutil.copyfileobj(source, sink)
+            while chunk := source.read(_EXTRACT_CHUNK_BYTES):
+                written += len(chunk)
+                if written > entry.file_size:
+                    raise BundleGateError(
+                        f'archive member {entry.filename!r} decompresses to more than its '
+                        f'declared {entry.file_size} bytes',
+                        gate=5,
+                    )
+                sink.write(chunk)
 
     # ------------------------------------------------------------------ gate 6
 

laborious/utils/repository/model_repository.py

@@ -139,9 +139,9 @@ class MLFlowRepository(SientiaMonitoring):
         """
         run_info = mlflow.get_run(run_id)
         if prediction:
-            model_uri = run_info.info.artifact_uri + '/prediction_model'
+            model_uri = run_info.info.artifact_uri + '/prediction_model'  # type: ignore[operator]
         else:
-            model_uri = run_info.info.artifact_uri + '/data_model'
+            model_uri = run_info.info.artifact_uri + '/data_model'  # type: ignore[operator]
         return model_uri
 
     def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str:
@@ -290,7 +290,11 @@ class MLFlowRepository(SientiaMonitoring):
     """
 
     async def dowload_artifacts(
-        self, model_name: str, metadata: dict[str, Any], artifact_path: str = 'data_model'
+        self,
+        model_name: str,
+        metadata: dict[str, Any],
+        artifact_path: str = 'data_model',
+        run_id: str | None = None,
     ) -> str:
         """
         Download artifacts from the latest production run of a model.
@@ -299,11 +303,14 @@ class MLFlowRepository(SientiaMonitoring):
             model_name (str): Registered model name.
             metadata (dict[str, Any]): Metadata used for structured logging.
             artifact_path (str): Relative path to artifacts within the run.
+            run_id (str | None): Run already resolved by the caller. When `None`,
+                the Production run is resolved here.
 
         Returns:
             str: Local filesystem path where artifacts are saved.
         """
-        run_id = self.get_model_run_id(model_name=model_name, stage='Production')
+        if run_id is None:
+            run_id = self.get_model_run_id(model_name=model_name, stage='Production')
         output_dir = f'{ARTIFACTS_PATH}/{model_name}'
 
         full_path = path.join(output_dir, artifact_path)
@@ -472,6 +479,7 @@ class MLFlowRepository(SientiaMonitoring):
         artifact_name: str,
         flavor: str,
         metadata: dict[str, Any] | None = None,
+        run_id: str | None = None,
     ) -> str | None:
         """
         Download only an artifact's `code` directory and put it on `sys.path`.
@@ -489,6 +497,8 @@ class MLFlowRepository(SientiaMonitoring):
             flavor (str): Flavor doing the load, for the log lines only.
             metadata (dict[str, Any] | None): Metadata used for structured
                 logging.
+            run_id (str | None): Run already resolved by the caller. When `None`,
+                the Production run is resolved here.
 
         Returns:
             str | None: Local directory the `code` directory was downloaded
@@ -500,7 +510,8 @@ class MLFlowRepository(SientiaMonitoring):
             Exception: Any error raised while resolving the Production run,
                 listing the artifact or downloading `code`.
         """
-        run_id = self.get_model_run_id(model_name=model_name, stage='Production')
+        if run_id is None:
+            run_id = self.get_model_run_id(model_name=model_name, stage='Production')
         remote_code_path = f'{artifact_name}/code'
 
         if not self.check_artifact_exists(run_id, remote_code_path, metadata):
@@ -623,7 +634,11 @@ class MLFlowRepository(SientiaMonitoring):
         return model, local_dir
 
     async def load_predict_model(
-        self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
+        self,
+        model_name: str,
+        metadata: dict[str, Any],
+        flavor: str = 'sklearn',
+        run_id: str | None = None,
     ) -> tuple[Any, str | None]:
         """
         Load a predictive model from the MLflow Model Registry.
@@ -632,6 +647,10 @@ class MLFlowRepository(SientiaMonitoring):
             model_name (str): The name of the model to download from the registry.
             metadata (dict[str, Any]): Metadata used for structured logging.
             flavor (str): Model flavor ('pyfunc', 'sklearn', 'joblib')
+            run_id (str | None): Run already resolved by the caller. When `None`,
+                the Production run is resolved where needed (only the `joblib`
+                branch downloads by run; the `pyfunc`/`sklearn` branches load
+                from the symbolic `models:/<name>/production` URI regardless).
 
         Returns:
             tuple[Any, str | None]: The loaded predictive model and the local
@@ -667,6 +686,7 @@ class MLFlowRepository(SientiaMonitoring):
                     model_name=model_name,
                     artifact_name='prediction_model',
                     metadata=metadata,
+                    run_id=run_id,
                 )
             elif flavor == 'pyfunc':
                 model, local_dir = mlflow.pyfunc.load_model(model_uri), None
@@ -681,7 +701,11 @@ class MLFlowRepository(SientiaMonitoring):
         return model, local_dir
 
     async def load_transform_model(
-        self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
+        self,
+        model_name: str,
+        metadata: dict[str, Any],
+        flavor: str = 'sklearn',
+        run_id: str | None = None,
     ) -> tuple[Any, str | None]:
         """
         Load the latest Production version of a transformation model.
@@ -693,6 +717,8 @@ class MLFlowRepository(SientiaMonitoring):
             model_name (str): The name of the model to download.
             metadata (dict[str, Any]): Metadata used for structured logging.
             flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib')
+            run_id (str | None): Run already resolved by the caller. When `None`,
+                the Production run is resolved here.
 
         Returns:
             tuple[Any, str | None]: The loaded model object, depending on the
@@ -719,7 +745,9 @@ class MLFlowRepository(SientiaMonitoring):
         if flavor not in ('sklearn', 'pyfunc', 'joblib'):
             raise ValueError(INVALID_FLAVOR_MESSAGE)
 
-        latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
+        latest_production_id = run_id or self.get_model_run_id(
+            model_name=model_name, stage='Production'
+        )
 
         model_uri = (
             None
@@ -802,8 +830,10 @@ class MLFlowRepository(SientiaMonitoring):
 
         target = 'prediction_model' if model_type == 'predict' else 'data_model'
 
+        run_id = self.get_model_run_id(model_name=model_name, stage='Production')
+
         code_artifact_dir = await self.download_artifact_code_dir(
-            model_name, target, flavor, metadata
+            model_name, target, flavor, metadata, run_id=run_id
         )
 
         if load_wrapper:
@@ -811,7 +841,9 @@ class MLFlowRepository(SientiaMonitoring):
             self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}')
 
             try:
-                artifact_path = await self.dowload_artifacts(model_name, metadata, target)
+                artifact_path = await self.dowload_artifacts(
+                    model_name, metadata, target, run_id=run_id
+                )
 
                 self.info(
                     f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
@@ -832,9 +864,13 @@ class MLFlowRepository(SientiaMonitoring):
             return model, code_artifact_dir or artifact_path
 
         if model_type == 'predict':
-            model, local_dir = await self.load_predict_model(model_name, metadata, flavor)
+            model, local_dir = await self.load_predict_model(
+                model_name, metadata, flavor, run_id=run_id
+            )
         else:
-            model, local_dir = await self.load_transform_model(model_name, metadata, flavor)
+            model, local_dir = await self.load_transform_model(
+                model_name, metadata, flavor, run_id=run_id
+            )
 
         return model, code_artifact_dir or local_dir
 

laborious/workflows/import_model.py

@@ -26,6 +26,7 @@ with workflow.unsafe.imports_passed_through():
     from typing import Any
 
     from sientia_do.temporal.policies import retry_policy
+    from temporalio.exceptions import is_cancelled_exception
 
     from laborious.activities.activities import Activities
     from laborious.activities.model_import import (
@@ -309,6 +310,8 @@ class ImportModel:
                 start_to_close_timeout=STATUS_WRITE_TIMEOUT,
             )
         except BaseException as error:  # NOSONAR - a lost hint must not fail a working import
+            if is_cancelled_exception(error):
+                raise
             await self._report_status_write_failure(base, error)
 
     async def _record_terminal(
@@ -344,6 +347,8 @@ class ImportModel:
         except (
             BaseException
         ) as error:  # NOSONAR - the outcome stands even when it cannot be written
+            if is_cancelled_exception(error):
+                raise
             await self._report_status_write_failure(
                 base,
                 error,
@@ -377,8 +382,9 @@ class ImportModel:
                 retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
                 start_to_close_timeout=STATUS_WRITE_TIMEOUT,
             )
-        except BaseException:  # noqa: S110  # NOSONAR - the report of a report is where this stops
-            pass
+        except BaseException as error:  # NOSONAR - the report of a report is where this stops
+            if is_cancelled_exception(error):
+                raise
 
     async def _report_status_write_failure(
         self,
@@ -394,10 +400,12 @@ class ImportModel:
                 retry_policy=retry_policy,
                 start_to_close_timeout=STATUS_WRITE_TIMEOUT,
             )
-        except BaseException:  # noqa: S110  # NOSONAR - the report of a report is where this stops
-            # Deliberately silent: the channel this was trying to reach is the thing that logs, so
-            # there is nowhere left to say it, and an import that worked must not fail over it.
-            pass
+        except BaseException as error:  # NOSONAR - the report of a report is where this stops
+            if is_cancelled_exception(error):
+                raise
+            # Deliberately silent otherwise: the channel this was trying to reach is the thing that
+            # logs, so there is nowhere left to say it, and an import that worked must not fail
+            # over it.
 
     # ----------------------------------------------------------- side-effect steps
 
@@ -418,12 +426,13 @@ class ImportModel:
                 retry_policy=retry_policy,
                 start_to_close_timeout=RETENTION_TIMEOUT,
             )
-        except BaseException:  # noqa: S110  # NOSONAR - retention is not a precondition of an import
-            # Dropped, not reported: the activity's own failure is already in the event history,
-            # and a bucket policy the importer cannot set is an operational problem rather than a
-            # bad import. Failing a valid model import over a lifecycle API would be the wrong
-            # trade.
-            pass
+        except BaseException as error:  # NOSONAR - retention is not a precondition of an import
+            if is_cancelled_exception(error):
+                raise
+            # Dropped, not reported otherwise: the activity's own failure is already in the event
+            # history, and a bucket policy the importer cannot set is an operational problem rather
+            # than a bad import. Failing a valid model import over a lifecycle API would be the
+            # wrong trade.
 
     @staticmethod
     async def _cleanup(base: dict[str, Any]) -> dict[str, Any]:
@@ -435,5 +444,7 @@ class ImportModel:
                 retry_policy=retry_policy,
                 start_to_close_timeout=CLEANUP_TIMEOUT,
             )
-        except BaseException:  # NOSONAR - cleanup never replaces the failure that brought us here
+        except BaseException as error:  # NOSONAR - cleanup never replaces the failure that brought us here
+            if is_cancelled_exception(error):
+                raise
             return {'cleaned': False}

tests/laborious/activities/test_model_import_bundle.py

@@ -72,24 +72,25 @@ def minio_client() -> MagicMock:
 def import_activity(tmp_path: Path, minio_client: MagicMock) -> ModelImport:
     repository = MagicMock()
     repository.minio_client = minio_client
-    with patch('laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()):
-        activity = ModelImport(
-            status_config={
-                'host': 'h',
-                'port': 5432,
-                'dbname': 'sientia-core-mlops-bff',
-                'user': 'u',
-                'password': 'p',
-                'schema': 'public',
-            },
-            import_config=import_config(tmp_path),
-            minio_repository=repository,
-            mlflow_repository=MagicMock(),
-            mongo_repository=AsyncMock(),
-            logger=MagicMock(),
-            notification_handler=MagicMock(),
-            metrics_controller=AsyncMock(),
-        )
+    activity = ModelImport(
+        status_config={
+            'host': 'h',
+            'port': 5432,
+            'dbname': 'sientia-core-mlops-bff',
+            'user': 'u',
+            'password': 'p',
+            'schema': 'public',
+        },
+        import_config=import_config(tmp_path),
+        minio_repository=repository,
+        mlflow_repository=MagicMock(),
+        mongo_repository=AsyncMock(),
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    # `status_engine` is created lazily on first access; pre-seed it with the recording double.
+    activity._status_engine = FakeStatusEngine()
     activity.emit_metric = AsyncMock()
     activity.send_notification_async = AsyncMock()
     return activity

tests/laborious/activities/test_model_import_provisioning.py

@@ -100,36 +100,37 @@ def import_activity(tmp_path: Path, mlflow_client: MagicMock, mongo: AsyncMock)
     minio_repository = MagicMock()
     mlflow_repository = MagicMock()
     mlflow_repository.client = mlflow_client
-    with patch('laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()):
-        activity = ModelImport(
-            status_config={
-                'host': 'h',
-                'port': 5432,
-                'dbname': 'sientia-core-mlops-bff',
-                'user': 'u',
-                'password': 'p',
-                'schema': 'public',
-            },
-            import_config={
-                'password_key': None,
-                'bucket': 'sientia',
-                'prefix': 'imported_models/',
-                'work_dir': str(tmp_path),
-                'max_object_bytes': 1024,
-                'max_entries': 10,
-                'max_uncompressed_bytes': 4096,
-                'max_compression_ratio': 200.0,
-                'models_collection': 'models',
-                'counters_collection': 'model_id_counters',
-                'retention_days': 7,
-            },
-            minio_repository=minio_repository,
-            mlflow_repository=mlflow_repository,
-            mongo_repository=mongo,
-            logger=MagicMock(),
-            notification_handler=MagicMock(),
-            metrics_controller=AsyncMock(),
-        )
+    activity = ModelImport(
+        status_config={
+            'host': 'h',
+            'port': 5432,
+            'dbname': 'sientia-core-mlops-bff',
+            'user': 'u',
+            'password': 'p',
+            'schema': 'public',
+        },
+        import_config={
+            'password_key': None,
+            'bucket': 'sientia',
+            'prefix': 'imported_models/',
+            'work_dir': str(tmp_path),
+            'max_object_bytes': 1024,
+            'max_entries': 10,
+            'max_uncompressed_bytes': 4096,
+            'max_compression_ratio': 200.0,
+            'models_collection': 'models',
+            'counters_collection': 'model_id_counters',
+            'retention_days': 7,
+        },
+        minio_repository=minio_repository,
+        mlflow_repository=mlflow_repository,
+        mongo_repository=mongo,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    # `status_engine` is created lazily on first access; pre-seed it with the recording double.
+    activity._status_engine = FakeStatusEngine()
     activity.emit_metric = AsyncMock()
     activity.send_notification_async = AsyncMock()
     return activity

tests/laborious/activities/test_model_import_status.py

@@ -164,17 +164,19 @@ def test_build_update_with_no_casts_or_extra_where_renders_a_bare_update():
 
 def build_activity(engine: FakeStatusEngine) -> ModelImport:
     """A `ModelImport` whose only connection is the recording engine double."""
-    with patch('laborious.activities.model_import.create_engine', return_value=engine):
-        activity = ModelImport(
-            status_config=STATUS_CONFIG,
-            import_config=IMPORT_CONFIG,
-            minio_repository=MagicMock(),
-            mlflow_repository=MagicMock(),
-            mongo_repository=AsyncMock(),
-            logger=MagicMock(),
-            notification_handler=MagicMock(),
-            metrics_controller=AsyncMock(),
-        )
+    activity = ModelImport(
+        status_config=STATUS_CONFIG,
+        import_config=IMPORT_CONFIG,
+        minio_repository=MagicMock(),
+        mlflow_repository=MagicMock(),
+        mongo_repository=AsyncMock(),
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    # `status_engine` is created lazily on first access; setting the private attribute directly
+    # pre-seeds it with the recording double instead of exercising the real `create_engine`.
+    activity._status_engine = engine
     activity.emit_metric = AsyncMock()
     activity.send_notification_async = AsyncMock()
     return activity
@@ -983,8 +985,62 @@ async def test_the_engine_is_disposed_on_shutdown(import_activity: ModelImport,
 @mark.asyncio
 async def test_close_is_a_no_op_without_an_engine(import_activity: ModelImport, engine):
     """A partially-constructed instance (or a second close) must not blow up on a missing engine."""
-    del import_activity.status_engine
+    import_activity._status_engine = None
 
     import_activity.close()
 
     assert engine.disposed is False
+
+
+# ------------------------------------------------------------- the lazy status engine
+
+
+def test_no_engine_is_created_when_status_engine_is_never_accessed():
+    """A worker pod that never runs an import activity never opens a pool against the BFF Postgres."""
+    with patch('laborious.activities.model_import.create_engine') as create_engine_mock:
+        activity = ModelImport(
+            status_config=STATUS_CONFIG,
+            import_config=IMPORT_CONFIG,
+            minio_repository=MagicMock(),
+            mlflow_repository=MagicMock(),
+            mongo_repository=AsyncMock(),
+            logger=MagicMock(),
+            notification_handler=MagicMock(),
+            metrics_controller=AsyncMock(),
+        )
+
+    create_engine_mock.assert_not_called()
+    assert activity._status_engine is None
+
+
+def test_first_access_creates_the_engine_and_later_accesses_reuse_it():
+    with patch(
+        'laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()
+    ) as create_engine_mock:
+        activity = ModelImport(
+            status_config=STATUS_CONFIG,
+            import_config=IMPORT_CONFIG,
+            minio_repository=MagicMock(),
+            mlflow_repository=MagicMock(),
+            mongo_repository=AsyncMock(),
+            logger=MagicMock(),
+            notification_handler=MagicMock(),
+            metrics_controller=AsyncMock(),
+        )
+
+        first = activity.status_engine
+        second = activity.status_engine
+
+    create_engine_mock.assert_called_once()
+    assert first is second
+
+
+@mark.asyncio
+async def test_close_never_accessed_the_engine_does_not_create_one(import_activity: ModelImport):
+    """`close()` must not force the lazy engine into existence just to dispose it."""
+    import_activity._status_engine = None
+
+    with patch('laborious.activities.model_import.create_engine') as create_engine_mock:
+        import_activity.close()
+
+    create_engine_mock.assert_not_called()

tests/laborious/utils/bundle/test_reader.py

@@ -437,7 +437,7 @@ def test_gate_5_wraps_an_os_error_during_unpacking(work_dir, reader, monkeypatch
     def _raise_os_error(*_args, **_kwargs):
         raise OSError('disk full')
 
-    monkeypatch.setattr('laborious.utils.bundle.reader.shutil.copyfileobj', _raise_os_error)
+    monkeypatch.setattr(zipfile.ZipExtFile, 'read', _raise_os_error)
 
     with pytest.raises(BundleGateError, match='could not be unpacked') as error:
         _open(reader, work_dir, payload)
@@ -463,6 +463,52 @@ def test_extract_member_rechecks_containment_on_the_resolved_path(work_dir, read
     assert error.value.gate == 5
 
 
+def test_extract_member_rejects_a_stream_larger_than_its_declared_size(work_dir, reader, monkeypatch):
+    """A decompressed stream longer than the entry's declared `file_size` is rejected mid-write,
+    instead of trusting the central-directory metadata gate 4 already checked in aggregate.
+
+    A single crafted DEFLATE chunk can make `ZipExtFile.read()` return far more bytes than
+    `file_size` declares before its own EOF/CRC bookkeeping catches up, so the stand-in here
+    forces exactly that response shape rather than fighting zipfile's own CRC validation."""
+    root = work_dir / 'root'
+    root.mkdir()
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, 'w') as archive:
+        archive.writestr('payload.bin', b'x' * 10)
+
+    call_count = {'n': 0}
+
+    def _oversized_read(self, *_args, **_kwargs):
+        call_count['n'] += 1
+        return b'x' * 1000 if call_count['n'] == 1 else b''
+
+    monkeypatch.setattr(zipfile.ZipExtFile, 'read', _oversized_read)
+
+    with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
+        entry = archive.infolist()[0]
+        with pytest.raises(BundleGateError, match='more than its declared') as error:
+            reader._extract_member(archive, entry, root, root.resolve())
+
+    assert error.value.gate == 5
+
+
+def test_extract_member_writes_a_well_formed_member_in_full(work_dir, reader):
+    """A member whose decompressed size matches its declared `file_size` is extracted byte for
+    byte, exactly as before the byte-counting check was added."""
+    root = work_dir / 'root'
+    root.mkdir()
+    payload = b'y' * 100
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, 'w') as archive:
+        archive.writestr('payload.bin', payload)
+
+    with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive:
+        entry = archive.infolist()[0]
+        reader._extract_member(archive, entry, root, root.resolve())
+
+    assert (root / 'payload.bin').read_bytes() == payload
+
+
 # ------------------------------------------------------------------- gate 6
 
 

tests/laborious/utils/repository/test_model_repository.py

@@ -557,12 +557,15 @@ async def test_download_model_invalid_model_type(mlflow_repository):
     'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')]
 )
 async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type, tmp_path):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
     mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path))
 
     result = await mlflow_repository.download_model('test_model', {}, model_type[0], 'pyfunc', True)
 
-    mlflow_repository.dowload_artifacts.assert_called_once_with('test_model', {}, model_type[1])
+    mlflow_repository.dowload_artifacts.assert_called_once_with(
+        'test_model', {}, model_type[1], run_id='run-123'
+    )
 
     mlflow.pyfunc.load_model.assert_called_once_with(
         mlflow_repository.dowload_artifacts.return_value
@@ -576,13 +579,16 @@ async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type
 
 @pytest.mark.asyncio
 async def test_download_model_predict(mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
     mlflow_repository.load_predict_model = AsyncMock(return_value=('model', 'local_dir'))
     mlflow_repository.load_transform_model = AsyncMock()
 
     result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', False)
 
-    mlflow_repository.load_predict_model.assert_called_once_with('test_model', {}, 'pyfunc')
+    mlflow_repository.load_predict_model.assert_called_once_with(
+        'test_model', {}, 'pyfunc', run_id='run-123'
+    )
     mlflow_repository.load_transform_model.assert_not_called()
 
     assert result == ('model', 'local_dir')
@@ -590,6 +596,7 @@ async def test_download_model_predict(mlflow_repository):
 
 @pytest.mark.asyncio
 async def test_download_model_transform(mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
     mlflow_repository.load_predict_model = AsyncMock()
     mlflow_repository.load_transform_model = AsyncMock(return_value=('model', 'local_dir'))
@@ -597,7 +604,9 @@ async def test_download_model_transform(mlflow_repository):
     result = await mlflow_repository.download_model('test_model', {}, 'transform', 'pyfunc', False)
 
     mlflow_repository.load_predict_model.assert_not_called()
-    mlflow_repository.load_transform_model.assert_called_once_with('test_model', {}, 'pyfunc')
+    mlflow_repository.load_transform_model.assert_called_once_with(
+        'test_model', {}, 'pyfunc', run_id='run-123'
+    )
 
     assert result == ('model', 'local_dir')
 
@@ -616,10 +625,11 @@ async def test_download_model_prepares_the_code_dir_first(
         calls.append('code')
         return 'code-artifact-dir'
 
-    async def record_load(*_a):
+    async def record_load(*_a, **_k):
         calls.append('load')
         return 'model', None
 
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(side_effect=record_code_dir)
     loader = AsyncMock(side_effect=record_load)
     mlflow_repository.load_predict_model = loader
@@ -628,7 +638,7 @@ async def test_download_model_prepares_the_code_dir_first(
     result = await mlflow_repository.download_model('test_model', {}, model_type, 'sklearn', False)
 
     mlflow_repository.download_artifact_code_dir.assert_called_once_with(
-        'test_model', artifact_name, 'sklearn', {}
+        'test_model', artifact_name, 'sklearn', {}, run_id='run-123'
     )
     assert calls == ['code', 'load']
     assert result == ('model', 'code-artifact-dir')
@@ -637,6 +647,7 @@ async def test_download_model_prepares_the_code_dir_first(
 @pytest.mark.asyncio
 async def test_download_model_wrapper_returns_the_code_artifact_dir(mlflow, mlflow_repository):
     """With a `code` directory downloaded, the wrapper branch returns it over its own download."""
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(return_value='code-artifact-dir')
     mlflow_repository.dowload_artifacts = AsyncMock(return_value='wrapper-dir')
 
@@ -647,6 +658,42 @@ async def test_download_model_wrapper_returns_the_code_artifact_dir(mlflow, mlfl
     assert artifact_dir == 'code-artifact-dir'
 
 
+@pytest.mark.asyncio
+async def test_download_model_resolves_run_id_once_non_wrapper(mlflow_repository):
+    """The `run_id` `download_model` resolves once is the same one every downstream call uses."""
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
+    mlflow_repository.load_predict_model = AsyncMock(return_value=('model', None))
+
+    await mlflow_repository.download_model('test_model', {}, 'predict', 'sklearn', False)
+
+    mlflow_repository.get_model_run_id.assert_called_once_with(model_name='test_model', stage='Production')
+    mlflow_repository.download_artifact_code_dir.assert_called_once_with(
+        'test_model', 'prediction_model', 'sklearn', {}, run_id='run-123'
+    )
+    mlflow_repository.load_predict_model.assert_called_once_with(
+        'test_model', {}, 'sklearn', run_id='run-123'
+    )
+
+
+@pytest.mark.asyncio
+async def test_download_model_resolves_run_id_once_wrapper(mlflow, mlflow_repository, tmp_path):
+    """Same single resolution, for the `load_wrapper=True` branch."""
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
+    mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path))
+
+    await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
+
+    mlflow_repository.get_model_run_id.assert_called_once_with(model_name='test_model', stage='Production')
+    mlflow_repository.download_artifact_code_dir.assert_called_once_with(
+        'test_model', 'prediction_model', 'pyfunc', {}, run_id='run-123'
+    )
+    mlflow_repository.dowload_artifacts.assert_called_once_with(
+        'test_model', {}, 'prediction_model', run_id='run-123'
+    )
+
+
 def test_detect_and_parse_datetime_index_empty(mlflow_repository):
     input_data = DataFrame()
     response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
@@ -1965,6 +2012,7 @@ async def test_load_transform_model_flavor_failure_propagates(
 @patch(JOBLIB_PATH)
 @pytest.mark.asyncio
 async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
     mlflow_repository.dowload_artifacts = AsyncMock()
     loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom')
@@ -1983,6 +2031,7 @@ async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, ml
 async def test_download_model_wrapper_failure_emits_error_metric(
     joblib_mock, mlflow, mlflow_repository
 ):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
     mlflow_repository.dowload_artifacts = AsyncMock()
     loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom')
@@ -2503,13 +2552,14 @@ async def test_download_model_delegates_every_flavor_to_the_loaders(
     mlflow_repository, model_type, loader
 ):
     """`download_model` owns no flavor branch: joblib delegates like the other two."""
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
     mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None)
     delegate = AsyncMock(return_value=('model', 'local-dir'))
     setattr(mlflow_repository, loader, delegate)
 
     result = await mlflow_repository.download_model('test_model', {}, model_type, 'joblib')
 
-    delegate.assert_called_once_with('test_model', {}, 'joblib')
+    delegate.assert_called_once_with('test_model', {}, 'joblib', run_id='run-123')
     assert result == ('model', 'local-dir')
 
 

tests/laborious/workflows/test_import_model.py

@@ -12,7 +12,7 @@ workflow is: a sequence of activity invocations with policies attached.
 import asyncio
 from datetime import timedelta
 from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from pytest import fixture, mark, raises
 from sientia_do.temporal.policies import retry_policy
@@ -725,6 +725,95 @@ async def test_reporting_the_report_is_where_it_stops(import_workflow: ImportMod
     assert result['run_id'] == RUN_ID
 
 
+# --------------------------------------------------- cancellation is never swallowed
+
+
+def cancelling_activities():
+    """Patch the workflow module's `workflow` object so every activity call raises cancellation."""
+    patcher = patch('laborious.workflows.import_model.workflow')
+    workflow_mock = patcher.start()
+    workflow_mock.execute_activity_method = AsyncMock(side_effect=asyncio.CancelledError())
+    return patcher
+
+
+BASE = {'metadata': {}, 'import_run_id': 77, 'workflow_id': WORKFLOW_ID}
+
+
+@mark.asyncio
+async def test_a_cancelled_hint_propagates_instead_of_being_reported(import_workflow: ImportModel):
+    patcher = cancelling_activities()
+    try:
+        with raises(asyncio.CancelledError):
+            await import_workflow._record_hint(
+                BASE, Activities.record_import_status, {'step': ImportStep.DOWNLOAD.value}
+            )
+    finally:
+        patcher.stop()
+
+
+@mark.asyncio
+async def test_a_cancelled_terminal_write_propagates_instead_of_being_reported(
+    import_workflow: ImportModel,
+):
+    patcher = cancelling_activities()
+    try:
+        with raises(asyncio.CancelledError):
+            await import_workflow._record_terminal(
+                BASE,
+                succeeded=True,
+                step=ImportStep.MODEL_DOCUMENT.value,
+                gate=None,
+                source={},
+                cleanup_failed=False,
+            )
+    finally:
+        patcher.stop()
+
+
+@mark.asyncio
+async def test_a_cancelled_pipeline_failure_report_propagates(import_workflow: ImportModel):
+    patcher = cancelling_activities()
+    try:
+        with raises(asyncio.CancelledError):
+            await import_workflow._report_failure(
+                BASE, RuntimeError('boom'), step=ImportStep.REGISTRATION.value, gate=None
+            )
+    finally:
+        patcher.stop()
+
+
+@mark.asyncio
+async def test_a_cancelled_status_write_failure_report_propagates(import_workflow: ImportModel):
+    patcher = cancelling_activities()
+    try:
+        with raises(asyncio.CancelledError):
+            await import_workflow._report_status_write_failure(BASE, RuntimeError('db down'))
+    finally:
+        patcher.stop()
+
+
+@mark.asyncio
+async def test_a_cancelled_retention_check_propagates(import_workflow: ImportModel):
+    patcher = cancelling_activities()
+    try:
+        with raises(asyncio.CancelledError):
+            await import_workflow._ensure_retention(BASE, valid_input())
+    finally:
+        patcher.stop()
+
+
+@mark.asyncio
+async def test_a_cancelled_cleanup_propagates_instead_of_returning_not_cleaned(
+    import_workflow: ImportModel,
+):
+    patcher = cancelling_activities()
+    try:
+        with raises(asyncio.CancelledError):
+            await import_workflow._cleanup(BASE)
+    finally:
+        patcher.stop()
+
+
 # --------------------------------------------------------------- the failure report