digitalkin 0.3.5.dev12__tar.gz → 0.3.5.dev13__tar.gz

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.5.dev12
+Version: 0.3.5.dev13
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -454,7 +454,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: ag-ui-protocol>=0.1.14
 Requires-Dist: agentic-mesh-protocol==0.2.3
-Requires-Dist: anyio==4.12.1
+Requires-Dist: anyio==4.13.0
 Requires-Dist: grpcio-health-checking==1.78.0
 Requires-Dist: grpcio-reflection==1.78.0
 Requires-Dist: grpcio-status==1.78.0
@@ -463,7 +463,7 @@ Provides-Extra: profiling
 Requires-Dist: asyncio-inspector==0.1.0; extra == "profiling"
 Requires-Dist: pyinstrument==5.1.2; extra == "profiling"
 Requires-Dist: viztracer==1.1.1; extra == "profiling"
-Requires-Dist: yappi==1.7.3; extra == "profiling"
+Requires-Dist: yappi==1.7.6; extra == "profiling"
 Provides-Extra: taskiq
 Requires-Dist: rstream==1.0.0; extra == "taskiq"
 Requires-Dist: taskiq-aio-pika==0.6.0; extra == "taskiq"

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/examples/modules/archetype_with_tools_module.py

@@ -5,10 +5,11 @@ from typing import Any, ClassVar, Literal
 
 from pydantic import BaseModel, Field
 
-from digitalkin.models.grpc_servers.models import ClientConfig, SecurityMode, ServerMode
+from digitalkin.models.grpc_servers.models import ClientConfig
 from digitalkin.models.module.module_context import ModuleContext
 from digitalkin.models.module.setup_types import SetupModel
 from digitalkin.models.module.tool_reference import ToolReference
+from digitalkin.models.settings.utils.channel import ControlFlow, SecurityMode
 from digitalkin.modules._base_module import BaseModule  # noqa: PLC2701
 from digitalkin.services.services_models import ServicesStrategy
 
@@ -98,7 +99,7 @@ class ArchetypeSecret(BaseModel):
 client_config = ClientConfig(
     host="[::]",
     port=50152,
-    mode=ServerMode.ASYNC,
+    mode=ControlFlow.ASYNC,
     security=SecurityMode.INSECURE,
     credentials=None,
 )

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/pyproject.toml

@@ -29,20 +29,20 @@
     dependencies = [
         "ag-ui-protocol>=0.1.14",
         "agentic-mesh-protocol==0.2.3",
-        "anyio==4.12.1",
+        "anyio==4.13.0",
         "grpcio-health-checking==1.78.0",
         "grpcio-reflection==1.78.0",
         "grpcio-status==1.78.0",
         "pydantic==2.12.5",
-]
-    version = "0.3.5.dev12"
+    ]
+    version = "0.3.5.dev13"
 
     [project.optional-dependencies]
         profiling = [
             "asyncio-inspector==0.1.0",
             "pyinstrument==5.1.2",
             "viztracer==1.1.1",
-            "yappi==1.7.3",
+            "yappi==1.7.6",
         ]
         taskiq = [
             "rstream==1.0.0",
@@ -61,13 +61,13 @@
 
 [dependency-groups]
     dev = [
-        "build==1.4.0",
+        "build==1.4.2",
         "bump-my-version==1.2.7",
-        "cryptography==46.0.5",
+        "cryptography==46.0.6",
         "mypy==1.19.1",
         "pre-commit==4.5.1",
         "pyright==1.1.408",
-        "ruff==0.15.5",
+        "ruff==0.15.8",
         "twine==6.2.0",
         "types-grpcio-health-checking==1.0.0.20250506",
         "types-grpcio-reflection==1.0.0.20250506",
@@ -87,17 +87,17 @@
         "mkdocs-git-revision-date-localized-plugin==1.5.1",
         "mkdocs-glightbox==0.5.2",
         "mkdocs-include-markdown-plugin==7.2.1",
-        "mkdocs-literate-nav==0.6.2",
+        "mkdocs-literate-nav==0.6.3",
         "mkdocs-llmstxt==0.5.0",
-        "mkdocs-material[imaging]==9.7.5",
+        "mkdocs-material[imaging]==9.7.6",
         "mkdocs-minify-plugin==0.8.0",
         "mkdocs-open-in-new-tab==1.0.8",
         "mkdocs-redirects==1.2.2",
-        "mkdocs-section-index==0.3.10",
+        "mkdocs-section-index==0.3.11",
         "mkdocs==1.6.1",
         "mkdocstrings-python==2.0.3",
         "mkdocstrings==1.0.3",
-        "tomli==2.4.0",
+        "tomli==2.4.1",
     ]
     tests = [
         "freezegun==1.5.5",
@@ -105,7 +105,7 @@
         "hdrhistogram==0.10.3",
         "psutil==7.2.2",
         "pytest-asyncio==1.3.0",
-        "pytest-cov==7.0.0",
+        "pytest-cov==7.1.0",
         "pytest-html==4.2.0",
         "pytest-json-report==1.5.0",
         "pytest-timeout==2.4.0",

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.5.dev12"
+    __version__ = "0.3.5.dev13"

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/src/digitalkin/community/agno/agno_adapter.py

@@ -86,7 +86,7 @@ class AgnoStreamAdapter:
 
         event_type = agno_event.event
 
-        logger.info("[DK STREAM-DEBUG => agno_adapter] Converting Agno event type: %s", event_type)
+        logger.debug("Converting Agno event: %s", event_type)
 
         timestamp = getattr(agno_event, "timestamp", None)
 
@@ -97,7 +97,7 @@ class AgnoStreamAdapter:
 
             # Skip duplicate RunStarted for the same run_id
             if run_id and run_id == self._active_run_id:
-                logger.info("[DK STREAM-DEBUG => agno_adapter] Skipping duplicate RunStarted for run_id=%s", run_id)
+                logger.debug("Skipping duplicate RunStarted for run_id=%s", run_id)
                 return []
 
             self._active_run_id = run_id
@@ -119,7 +119,7 @@ class AgnoStreamAdapter:
 
             # Skip if this run was already completed and not reopened
             if run_id and run_id in self._completed_run_ids and run_id != self._active_run_id:
-                logger.info("[DK STREAM-DEBUG => agno_adapter] Skipping duplicate RunCompleted for run_id=%s", run_id)
+                logger.debug("Skipping duplicate RunCompleted for run_id=%s", run_id)
                 return []
 
             events: list[BaseAgentRunEvent] = []
@@ -169,10 +169,7 @@ class AgnoStreamAdapter:
 
             self._current_reasoning_id = str(uuid.uuid4())
             self._reasoning_active = True
-            logger.info(
-                "[DK STREAM-DEBUG => agno_adapter] Reasoning started (explicit), id=%s",
-                self._current_reasoning_id,
-            )
+            logger.debug("Reasoning started, id=%s", self._current_reasoning_id)
             events.append(
                 ReasoningStartedEvent(
                     event=AgentRunEvent.REASONING_STARTED,
@@ -208,7 +205,7 @@ class AgnoStreamAdapter:
             ]
 
         if event_type == RunEvent.reasoning_completed:
-            logger.info("[DK STREAM-DEBUG => agno_adapter] Reasoning completed (explicit)")
+            logger.debug("Reasoning completed")
             return self._close_reasoning(timestamp)
 
         # ── Tool Call Events ─────────────────────────────────────────────
@@ -218,7 +215,7 @@ class AgnoStreamAdapter:
 
             # Close both reasoning and content before tool calls
             if self._reasoning_active:
-                logger.info("[DK STREAM-DEBUG => agno_adapter] Reasoning auto-completed (tool call started)")
+                logger.debug("Reasoning auto-completed (tool call started)")
                 events.extend(self._close_reasoning(timestamp))
             if self._content_active:
                 events.extend(self._close_content(timestamp))
@@ -275,10 +272,7 @@ class AgnoStreamAdapter:
 
             # Skip if already closed by ToolCallCompleted (Agno emits both for the same tool_call_id)
             if tool_call_id and tool_call_id in self._closed_tool_call_ids:
-                logger.info(
-                    "[DK STREAM-DEBUG => agno_adapter] Skipping duplicate ToolCallError for already-closed tool %s",
-                    tool_call_id,
-                )
+                logger.debug("Skipping duplicate ToolCallError for tool %s", tool_call_id)
                 return []
 
             tool_info = None
@@ -315,10 +309,10 @@ class AgnoStreamAdapter:
         """
         events: list[BaseAgentRunEvent] = []
         if self._content_active:
-            logger.info("[DK STREAM-DEBUG => agno_adapter] Flushing: closing active content")
+            logger.debug("Flushing active content sequence")
             events.extend(self._close_content(None))
         if self._reasoning_active:
-            logger.info("[DK STREAM-DEBUG => agno_adapter] Flushing: closing active reasoning")
+            logger.debug("Flushing active reasoning sequence")
             events.extend(self._close_reasoning(None))
         return events
 
@@ -394,7 +388,7 @@ class AgnoStreamAdapter:
 
         # Edge case: neither reasoning_content nor content
         if reasoning_content is None and content is None:
-            logger.debug("[DK STREAM-DEBUG => agno_adapter] run_content with no content, skipping")
+            logger.debug("run_content with no content, skipping")
 
         return events
 
@@ -420,10 +414,7 @@ class AgnoStreamAdapter:
         # Auto-open reasoning on first chunk
         if not self._reasoning_active:
             self._current_reasoning_id = str(uuid.uuid4())
-            logger.info(
-                "[DK STREAM-DEBUG => agno_adapter] Reasoning auto-started from run_content, id=%s",
-                self._current_reasoning_id,
-            )
+            logger.debug("Reasoning auto-started, id=%s", self._current_reasoning_id)
             events.append(
                 ReasoningStartedEvent(
                     event=AgentRunEvent.REASONING_STARTED,
@@ -434,10 +425,6 @@ class AgnoStreamAdapter:
             )
             self._reasoning_active = True
 
-        logger.info(
-            "[DK STREAM-DEBUG => agno_adapter] Reasoning content delta (len=%d)",
-            len(reasoning_content),
-        )
         events.append(
             ReasoningContentDeltaEvent(
                 event=AgentRunEvent.REASONING_CONTENT_DELTA,
@@ -466,7 +453,7 @@ class AgnoStreamAdapter:
         # Non-empty string → text data
         # Close reasoning if transitioning from reasoning to content
         if self._reasoning_active:
-            logger.info("[DK STREAM-DEBUG => agno_adapter] Reasoning auto-completed (text content arrived)")
+            logger.debug("Reasoning auto-completed (text content arrived)")
             events.extend(self._close_reasoning(timestamp))
 
         # Auto-open text message on first chunk

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/src/digitalkin/core/job_manager/base_job_manager.py

@@ -130,6 +130,12 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
         required for the job manager to function.
         """
 
+    async def stop(self) -> None:
+        """Stop the job manager and clean up resources.
+
+        Default no-op. Subclasses with external connections override.
+        """
+
     @staticmethod
     async def job_specific_callback(
         callback: Callable[[str, DataModel | ModuleCodeModel], Coroutine[Any, Any, None]],

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/src/digitalkin/core/job_manager/single_job_manager.py

@@ -285,10 +285,7 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
                 if session.cancelled:
                     break
 
-        try:
-            yield _stream()
-        finally:
-            session.close_stream()
+        yield _stream()
 
     async def create_module_instance_job(
         self,

{digitalkin-0.3.5.dev12 → digitalkin-0.3.5.dev13}/src/digitalkin/core/job_manager/taskiq_broker.py

@@ -1,18 +1,20 @@
 """Taskiq broker & RSTREAM producer for the job manager."""
 
 import asyncio
-import json
 import logging
 import os
 import pickle
+import ssl
 from typing import Any
 
 from rstream import Producer
 from rstream.exceptions import PreconditionFailed
 from taskiq import Context, TaskiqDepends, TaskiqMessage
 from taskiq.abc.formatter import TaskiqFormatter
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.compat import model_validate
 from taskiq.message import BrokerMessage
+from taskiq.result import TaskiqResult
 from taskiq_aio_pika import AioPikaBroker
 
 from digitalkin.core.common import ModuleFactory
@@ -61,27 +63,79 @@ class PickleFormatter(TaskiqFormatter):
     def loads(self, message: bytes) -> TaskiqMessage:  # Required by TaskiqFormatter interface # noqa: PLR6301
         """Recreate Python object from bytes.
 
+        Non-pickle messages (e.g. raw JSON left in the queue by other producers)
+        are logged and converted to a no-op ``TaskiqMessage`` so that Taskiq
+        acknowledges (consumes) them instead of nack-ing and re-delivering in a loop.
+
         Args:
             message: Broker message from bytes.
 
         Returns:
             message with TaskIQ format
         """
-        json_str = pickle.loads(  # noqa: S301
-            message
-        )  # Pickle: required for Taskiq deserialization (internal broker messages only)
+        try:
+            json_str = pickle.loads(  # noqa: S301
+                message
+            )  # Pickle: required for Taskiq deserialization (internal broker messages only)
+        except Exception as e:
+            logger.warning(
+                "Discarding non-pickle message (size=%d, preview=%r): %s",
+                len(message),
+                message[:80],
+                e,
+            )
+            # Return a no-op message that Taskiq will ack and discard
+            # (no task named "__discarded__" exists, so Taskiq logs a warning and moves on)
+            return TaskiqMessage(
+                task_id="__discarded__",
+                task_name="__discarded__",
+                labels={"_discarded": "true"},
+                args=[],
+                kwargs={},
+            )
         return model_validate(TaskiqMessage, json_str)
 
 
+def _rstream_ssl_context() -> ssl.SSLContext | None:
+    """Create SSL context for RStream if TLS is enabled via RABBITMQ_RSTREAM_SSL=true.
+
+    Returns:
+        SSL context if TLS is enabled, None otherwise.
+    """
+    if os.environ.get("RABBITMQ_RSTREAM_SSL", "").lower() not in {"true", "1", "yes"}:
+        return None
+    ctx = ssl.create_default_context()
+    # Allow self-signed certs in staging if RABBITMQ_RSTREAM_SSL_VERIFY=false
+    if os.environ.get("RABBITMQ_RSTREAM_SSL_VERIFY", "true").lower() in {"false", "0", "no"}:
+        ctx.check_hostname = False
+        ctx.verify_mode = ssl.CERT_NONE
+    return ctx
+
+
 class TaskiqBrokerConfig:
     """Configuration and lifecycle management for Taskiq broker and RStream producer."""
 
     STREAM = "taskiq_data"
     STREAM_RETENTION = 200_000
 
+    @staticmethod
+    async def _on_producer_closed(reason: Any) -> None:
+        """Log RStream producer connection closure for diagnostics.
+
+        Args:
+            reason: Connection close reason from rstream.
+        """
+        logger.error("RStream producer connection closed: %s", reason)
+
     @staticmethod
     def define_producer() -> Producer:
-        """Create RStream producer from environment variables.
+        """Create RStream producer with tuned settings for sustained throughput.
+
+        Tuning:
+        - ``default_batch_publishing_delay``: Flush batches every 100ms (default 3s)
+          for lower streaming latency during long-running tasks.
+        - ``default_context_switch_value``: Yield to the event loop every 100 messages
+          (default 1000) to keep concurrent coroutines responsive under heavy output.
 
         Returns:
             Producer connected to RabbitMQ.
@@ -91,12 +145,22 @@ class TaskiqBrokerConfig:
         username = os.environ.get("RABBITMQ_RSTREAM_USERNAME", "guest")
         password = os.environ.get("RABBITMQ_RSTREAM_PASSWORD", "guest")
 
-        logger.info("Connection to RabbitMQ: %s:%s.", host, port)
-        return Producer(host=host, port=int(port), username=username, password=password)
+        logger.info("RStream producer connecting to %s:%s", host, port)
+        return Producer(
+            host=host,
+            port=int(port),
+            username=username,
+            password=password,
+            ssl_context=_rstream_ssl_context(),
+            default_batch_publishing_delay=float(os.environ.get("DIGITALKIN_RSTREAM_BATCH_DELAY", "0.1")),
+            default_context_switch_value=int(os.environ.get("DIGITALKIN_RSTREAM_CONTEXT_SWITCH", "100")),
+            connection_name="digitalkin_producer",
+            on_close_handler=TaskiqBrokerConfig._on_producer_closed,
+        )
 
     @staticmethod
     def define_broker() -> AioPikaBroker:
-        """Create AioPikaBroker from environment variables.
+        """Create AioPikaBroker with tuned QoS for worker prefetch control.
 
         Returns:
             Broker connected to RabbitMQ with custom formatter.
@@ -105,12 +169,19 @@ class TaskiqBrokerConfig:
         port = os.environ.get("RABBITMQ_BROKER_PORT", "5672")
         username = os.environ.get("RABBITMQ_BROKER_USERNAME", "guest")
         password = os.environ.get("RABBITMQ_BROKER_PASSWORD", "guest")
+        scheme = os.environ.get("RABBITMQ_BROKER_SCHEME", "amqp")
 
         broker = AioPikaBroker(
-            f"amqp://{username}:{password}@{host}:{port}",
+            f"{scheme}://{username}:{password}@{host}:{port}",
+            qos=int(os.environ.get("DIGITALKIN_TASKIQ_PREFETCH", "10")),
             startup=[TaskiqBrokerConfig.init_rstream],
         )
         broker.formatter = PickleFormatter()
+        redis_url = os.environ.get("DIGITALKIN_TASKIQ_RESULT_BACKEND_URL")
+        if redis_url:
+            from taskiq_redis import RedisAsyncResultBackend
+
+            broker.with_result_backend(RedisAsyncResultBackend(redis_url))
         return broker
 
     @staticmethod
@@ -147,17 +218,78 @@ class TaskiqBrokerConfig:
     async def send_message_to_stream(job_id: str, output_data: DataModel | ModuleCodeModel) -> None:
         """Add a message frame to the RStream.
 
+        Uses Pydantic's Rust-based model_dump_json() and direct string embedding
+        to avoid the overhead of model_dump() → dict → json.dumps() → encode().
+
         Args:
             job_id: ID of the job that sent the message.
             output_data: Message body as a OutputModelT or error / stream_code.
         """
-        body = json.dumps({"job_id": job_id, "output_data": output_data.model_dump(mode="json")}).encode("utf-8")
+        # job_id is always a UUID (hex + hyphens), safe to embed without escaping
+        output_json = output_data.model_dump_json()
+        body = f'{{"job_id":"{job_id}","output_data":{output_json}}}'.encode()
         await RSTREAM_PRODUCER.send(stream=TaskiqBrokerConfig.STREAM, message=body)
 
 
+class TaskiqLifecycleMiddleware(TaskiqMiddleware):
+    """Lifecycle middleware for structured logging and safety-net EndOfStreamOutput."""
+
+    async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:  # noqa: PLR6301
+        """Log task start.
+
+        Returns:
+            The unmodified message.
+        """
+        logger.info("Taskiq task starting: %s (task_name=%s)", message.task_id, message.task_name)
+        return message
+
+    async def post_execute(self, message: TaskiqMessage, result: TaskiqResult) -> None:  # noqa: PLR6301
+        """Log task completion."""
+        log_fn = logger.info if not result.is_err else logger.error
+        log_fn(
+            "Taskiq task finished: %s (task_name=%s, is_err=%s, exec_time=%.3fs)",
+            message.task_id,
+            message.task_name,
+            result.is_err,
+            result.execution_time,
+        )
+
+    async def on_error(  # noqa: PLR6301
+        self,
+        message: TaskiqMessage,
+        result: TaskiqResult,  # noqa: ARG002
+        exception: BaseException,
+    ) -> None:
+        """Safety net: send EndOfStreamOutput if worker task failed to."""
+        logger.error("Taskiq task error: %s (task_name=%s, error=%s)", message.task_id, message.task_name, exception)
+        try:
+            await TaskiqBrokerConfig.send_message_to_stream(
+                message.task_id,
+                ModuleCodeModel(code="WorkerCrash", short_description="Middleware safety net", message=str(exception)),
+            )
+            await TaskiqBrokerConfig.send_message_to_stream(
+                message.task_id,
+                DataModel(root=EndOfStreamOutput()),
+            )
+        except Exception:
+            logger.exception("Middleware safety net failed for %s", message.task_id)
+
+
 # Module-level globals required by Taskiq framework (decorator needs broker at import time)
 RSTREAM_PRODUCER = TaskiqBrokerConfig.define_producer()
 TASKIQ_BROKER = TaskiqBrokerConfig.define_broker()
+TASKIQ_BROKER.add_middlewares(TaskiqLifecycleMiddleware())
+
+
+@TASKIQ_BROKER.task(task_name="__discarded__")
+async def _discarded_message() -> None:  # noqa: RUF029
+    """No-op sink for poison messages consumed by PickleFormatter.
+
+    Taskiq's receiver early-returns without acking when a task name is unknown,
+    so we register this dummy task to ensure the message is executed (no-op),
+    acked, and removed from the queue.
+    """
+    logger.debug("Poison message acknowledged and discarded")
 
 
 @TASKIQ_BROKER.task
@@ -170,6 +302,7 @@ async def run_start_module(
     input_data: dict,
     setup_data: dict,
     request_metadata: dict[str, str] | None = None,
+    registry_config: dict[str, Any] | None = None,
     context: Context = TaskiqDepends(),
 ) -> None:
     """TaskIQ task allowing a module to compute in the background asynchronously.
@@ -183,9 +316,17 @@ async def run_start_module(
         input_data: dict,
         setup_data: dict,
         request_metadata: gRPC request metadata (headers) to forward to the module.
+        registry_config: Registry config (client_config) forwarded from the main process.
         context: Allow TaskIQ context access
     """
     logger.info("Starting module with services_mode: %s", services_mode)
+
+    # Restore registry config lost during pickle (worker re-imports class without runtime mutations)
+    if registry_config is not None:
+        if "services_config_params" not in module_class.__dict__:
+            module_class.services_config_params = dict(module_class.services_config_params)
+        module_class.services_config_params["registry"] = registry_config
+
     services_config = ServicesConfig(
         services_config_strategies=module_class.services_config_strategies,
         services_config_params=module_class.services_config_params,
@@ -219,6 +360,17 @@ async def run_start_module(
             setup_model = await module_class.create_setup_model(setup_data)
         except Exception as e:
             logger.error("Failed to reconstruct models for job %s: %s", job_id, e, exc_info=True)
+            try:
+                await callback(
+                    ModuleCodeModel(
+                        code="ValidationError",
+                        short_description="Model reconstruction failed",
+                        message=str(e),
+                    )
+                )
+                await callback(DataModel(root=EndOfStreamOutput()))
+            except Exception:
+                logger.exception("Failed to send error to stream for job %s", job_id)
             raise
 
         supervisor_task = await executor.execute_task(
@@ -236,8 +388,19 @@ async def run_start_module(
         # Wait for the supervisor task to complete
         await supervisor_task
         logger.info("Module task %s completed", job_id)
-    except Exception:
+    except Exception as e:
         logger.exception("Error running module %s", job_id)
+        try:
+            await callback(
+                ModuleCodeModel(
+                    code="WorkerError",
+                    short_description="Worker execution failed",
+                    message=str(e),
+                )
+            )
+            await callback(DataModel(root=EndOfStreamOutput()))
+        except Exception:
+            logger.exception("Failed to send error to stream for job %s", job_id)
         raise
     finally:
         # Cleanup via module context
@@ -256,6 +419,7 @@ async def run_config_module(
     services_mode: ServicesMode,
     config_setup_data: dict,
     request_metadata: dict[str, str] | None = None,
+    registry_config: dict[str, Any] | None = None,
     context: Context = TaskiqDepends(),
 ) -> None:
     """TaskIQ task allowing a module to compute in the background asynchronously.
@@ -268,9 +432,17 @@ async def run_config_module(
         services_mode: ServicesMode,
         config_setup_data: dict,
         request_metadata: gRPC request metadata (headers) to forward to the module.
+        registry_config: Registry config (client_config) forwarded from the main process.
         context: Allow TaskIQ context access
     """
     logger.info("Starting config module with services_mode: %s", services_mode)
+
+    # Restore registry config lost during pickle (worker re-imports class without runtime mutations)
+    if registry_config is not None:
+        if "services_config_params" not in module_class.__dict__:
+            module_class.services_config_params = dict(module_class.services_config_params)
+        module_class.services_config_params["registry"] = registry_config
+
     services_config = ServicesConfig(
         services_config_strategies=module_class.services_config_strategies,
         services_config_params=module_class.services_config_params,
@@ -293,7 +465,22 @@ async def run_config_module(
         session = TaskSession(job_id, mission_id, module)
 
         # Create and run the config setup task with TaskExecutor
-        setup_model = module_class.create_config_setup_model(config_setup_data)
+        try:
+            setup_model = module_class.create_config_setup_model(config_setup_data)
+        except Exception as e:
+            logger.error("Failed to reconstruct config setup model for job %s: %s", job_id, e, exc_info=True)
+            try:
+                await callback(
+                    ModuleCodeModel(
+                        code="ValidationError",
+                        short_description="Config setup model reconstruction failed",
+                        message=str(e),
+                    )
+                )
+                await callback(DataModel(root=EndOfStreamOutput()))
+            except Exception:
+                logger.exception("Failed to send error to stream for job %s", job_id)
+            raise
 
         supervisor_task = await executor.execute_task(
             task_id=job_id,
@@ -305,8 +492,19 @@ async def run_config_module(
         # Wait for the supervisor task to complete
         await supervisor_task
         logger.info("Config module task %s completed", job_id)
-    except Exception:
+    except Exception as e:
         logger.exception("Error running config module %s", job_id)
+        try:
+            await callback(
+                ModuleCodeModel(
+                    code="WorkerError",
+                    short_description="Config worker execution failed",
+                    message=str(e),
+                )
+            )
+            await callback(DataModel(root=EndOfStreamOutput()))
+        except Exception:
+            logger.exception("Failed to send error to stream for job %s", job_id)
         raise
     finally:
         # Cleanup via module context