digitalkin 0.3.2.dev17__tar.gz → 0.3.2.dev19__tar.gz

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.2.dev17
+Version: 0.3.2.dev19
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/pyproject.toml

@@ -12,7 +12,7 @@
 
     keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
 
-    version = "0.3.2.dev17"
+    version = "0.3.2.dev19"
     classifiers = [
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.2.dev17"
+    __version__ = "0.3.2.dev19"

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/src/digitalkin/grpc_servers/module_servicer.py

@@ -17,7 +17,7 @@ from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.grpc_servers.utils.exceptions import ServicerError
 from digitalkin.logger import logger
 from digitalkin.models.core.job_manager_models import JobManagerMode
-from digitalkin.models.module.module import ModuleStatus
+from digitalkin.models.module.module import ModuleCodeModel, ModuleStatus
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.registry import GrpcRegistry, RegistryStrategy
 from digitalkin.services.services_models import ServicesMode
@@ -159,7 +159,32 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
 
         updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
-        logger.info("Setup updated", extra={"job_id": job_id})
+        logger.info("Setup response received", extra={"job_id": job_id})
+
+        # Check if response is an error
+        if isinstance(updated_setup_data, ModuleCodeModel):
+            logger.error(
+                "Config setup failed",
+                extra={"job_id": job_id, "code": updated_setup_data.code, "message": updated_setup_data.message},
+            )
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(updated_setup_data.message or "Config setup failed")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
+        if isinstance(updated_setup_data, dict) and "code" in updated_setup_data:
+            # ModuleCodeModel was serialized to dict
+            logger.error(
+                "Config setup failed",
+                extra={
+                    "job_id": job_id,
+                    "code": updated_setup_data["code"],
+                    "message": updated_setup_data.get("message"),
+                },
+            )
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(updated_setup_data.get("message") or "Config setup failed")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
         logger.debug("Updated setup data", extra={"job_id": job_id, "setup_data": updated_setup_data})
         setup_version.content = json_format.ParseDict(
             updated_setup_data,

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/src/digitalkin/models/module/module_context.py

@@ -8,7 +8,7 @@ from typing import Any
 from zoneinfo import ZoneInfo
 
 from digitalkin.logger import logger
-from digitalkin.models.module.tool_cache import ToolCache, ToolDefinition, ToolParameter
+from digitalkin.models.module.tool_cache import ToolCache, ToolDefinition, ToolModuleInfo, ToolParameter
 from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.communication.communication_strategy import CommunicationStrategy
 from digitalkin.services.cost.cost_strategy import CostStrategy
@@ -227,20 +227,20 @@ class ModuleContext:
             llm_format=llm_format,
         )
 
-    async def create_openai_style_tools(self, tool_name: str) -> list[dict[str, Any]]:
+    async def create_openai_style_tools(self, module_id: str) -> list[dict[str, Any]]:
         """Create OpenAI-style function calling schemas for a tool module.
 
         Uses tool cache (fast path) with registry fallback. Returns one schema
         per ToolDefinition (protocol) in the module.
 
         Args:
-            tool_name: Module ID to look up (checks cache first, then registry).
+            module_id: Module ID to look up (checks cache first, then registry).
 
         Returns:
             List of OpenAI-style tool schemas, one per protocol. Empty if not found.
         """
         tool_module_info = await self.tool_cache.get(
-            tool_name, registry=self.registry, communication=self.communication
+            module_id, registry=self.registry, communication=self.communication
         )
         if not tool_module_info:
             return []
@@ -299,13 +299,11 @@ class ModuleContext:
 
         communication = self.communication
         session = self.session
-        address = tool_module_info.address
-        port = tool_module_info.port
 
         result = []
         for tool_def in tool_module_info.tools:
             # Capture tool_def in closure via separate method
-            fn = ModuleContext._create_single_tool_function(communication, session, address, port, tool_def)
+            fn = ModuleContext._create_single_tool_function(communication, session, tool_module_info, tool_def)
             result.append((tool_def, fn))
 
         return result
@@ -314,8 +312,7 @@ class ModuleContext:
     def _create_single_tool_function(
         communication: CommunicationStrategy,
         session: Session,
-        address: str,
-        port: int,
+        tool_module_info: ToolModuleInfo,
         tool_def: ToolDefinition,
     ) -> Callable[..., AsyncGenerator[dict, None]]:
         """Create a single tool function for a specific protocol.
@@ -323,8 +320,7 @@ class ModuleContext:
         Args:
             communication: Communication strategy for gRPC calls.
             session: Current session with setup_id and mission_id.
-            address: Module address.
-            port: Module port.
+            tool_module_info: Tool module information containing address and port.
             tool_def: Tool definition with protocol name.
 
         Returns:
@@ -336,10 +332,10 @@ class ModuleContext:
             kwargs["protocol"] = protocol
             wrapped_input = {"root": kwargs}
             async for response in communication.call_module(
-                module_address=address,
-                module_port=port,
+                module_address=tool_module_info.address,
+                module_port=tool_module_info.port,
                 input_data=wrapped_input,
-                setup_id=session.setup_id,
+                setup_id=tool_module_info.setup_id,
                 mission_id=session.mission_id,
             ):
                 yield response

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/src/digitalkin/models/module/setup_types.py

@@ -29,65 +29,12 @@ SetupModelT = TypeVar("SetupModelT", bound="SetupModel")
 class SetupModel(BaseModel, Generic[SetupModelT]):
     """Base setup model with dynamic schema and tool cache support."""
 
+    model_config = ConfigDict(extra="allow")
     _clean_model_cache: ClassVar[dict[tuple[type, bool, bool], type]] = {}
-
-    def __init_subclass__(cls, **kwargs: Any) -> None:  # noqa: ANN401
-        """Inject hidden companion fields for ToolReference annotations.
-
-        Args:
-            **kwargs: Keyword arguments passed to parent.
-        """
-        super().__init_subclass__(**kwargs)
-        cls._inject_tool_cache_fields()
-
-    @classmethod
-    def _inject_tool_cache_fields(cls) -> None:
-        """Inject hidden companion fields for ToolReference annotations."""
-        annotations = getattr(cls, "__annotations__", {})
-        new_annotations: dict[str, Any] = {}
-
-        for field_name, annotation in annotations.items():
-            if cls._is_tool_reference_annotation(annotation):
-                cache_field_name = f"{field_name}_cache"
-                if cache_field_name not in annotations:
-                    # Check if it's a list type
-                    origin = get_origin(annotation)
-                    if origin is list:
-                        new_annotations[cache_field_name] = list[ToolModuleInfo]
-                        setattr(
-                            cls,
-                            cache_field_name,
-                            Field(default_factory=list, json_schema_extra={"hidden": True}),
-                        )
-                    else:
-                        new_annotations[cache_field_name] = ToolModuleInfo | None
-                        setattr(
-                            cls,
-                            cache_field_name,
-                            Field(default=None, json_schema_extra={"hidden": True}),
-                        )
-
-        if new_annotations:
-            cls.__annotations__ = {**annotations, **new_annotations}
-
-    @classmethod
-    def _is_tool_reference_annotation(cls, annotation: object) -> bool:
-        """Check if annotation is ToolReference or Optional[ToolReference].
-
-        Args:
-            annotation: Type annotation to check.
-
-        Returns:
-            True if annotation is or contains ToolReference.
-        """
-        origin = get_origin(annotation)
-        if origin is typing.Union or origin is types.UnionType:
-            return any(
-                arg is ToolReference or (isinstance(arg, type) and issubclass(arg, ToolReference))
-                for arg in get_args(annotation)
-                if arg is not type(None)
-            )
-        return annotation is ToolReference or (isinstance(annotation, type) and issubclass(annotation, ToolReference))
+    resolved_tools: dict[str, ToolModuleInfo] = Field(
+        default_factory=dict,
+        json_schema_extra={"hidden": True},
+    )
 
     @classmethod
     async def get_clean_model(
@@ -285,7 +232,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         if not has_changes:
             return model_cls
 
-        root_extra = cls.model_config.get("json_schema_extra", {})
+        root_extra = model_cls.model_config.get("json_schema_extra", {})
 
         return create_model(
             model_cls.__name__,
@@ -346,7 +293,12 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             communication: Communication service for module schemas.
         """
         logger.info("Starting resolve_tool_references")
-        await self._resolve_tool_references_recursive(self, registry, communication)
+        await self._resolve_tool_references_recursive(
+            self,
+            registry,
+            communication,
+            self.resolved_tools,
+        )
         logger.info("Finished resolve_tool_references")
 
     @classmethod
@@ -355,6 +307,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         model_instance: BaseModel,
         registry: "RegistryStrategy",
         communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Recursively resolve ToolReference fields in a model.
 
@@ -362,11 +315,18 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             model_instance: Model instance to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         for field_name, field_value in model_instance.__dict__.items():
             if field_value is None:
                 continue
-            await cls._resolve_field_value(field_name, field_value, registry, communication)
+            await cls._resolve_field_value(
+                field_name,
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
 
     @classmethod
     async def _resolve_field_value(
@@ -375,6 +335,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         field_value: "BaseModel | ToolReference | list | dict",
         registry: "RegistryStrategy",
         communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve a single field value based on its type.
 
@@ -383,15 +344,37 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             field_value: Value to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         if isinstance(field_value, ToolReference):
-            await cls._resolve_single_tool_reference(field_name, field_value, registry, communication)
+            await cls._resolve_single_tool_reference(
+                field_name,
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
         elif isinstance(field_value, BaseModel):
-            await cls._resolve_tool_references_recursive(field_value, registry, communication)
+            await cls._resolve_tool_references_recursive(
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
         elif isinstance(field_value, list):
-            await cls._resolve_list_items(field_value, registry, communication)
+            await cls._resolve_list_items(
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
         elif isinstance(field_value, dict):
-            await cls._resolve_dict_values(field_value, registry, communication)
+            await cls._resolve_dict_values(
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
 
     @classmethod
     async def _resolve_single_tool_reference(
@@ -400,6 +383,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         tool_ref: ToolReference,
         registry: "RegistryStrategy",
         communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve a single ToolReference.
 
@@ -408,17 +392,33 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             tool_ref: ToolReference to resolve.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         logger.info("Resolving ToolReference '%s' with setup_id='%s'", field_name, tool_ref.config.setup_id)
+
+        slug = tool_ref.slug
+        if slug:
+            cached = resolved_tools.get(slug)
+            if cached:
+                tool_ref._cached_info = cached  # noqa: SLF001
+                logger.info("ToolReference '%s' resolved from cache -> %s", field_name, cached)
+                return
+
         try:
-            await tool_ref.resolve(registry, communication)
+            info = await tool_ref.resolve(registry, communication)
+            if info and info.setup_id:
+                resolved_tools[info.setup_id] = info
             logger.info("Resolved ToolReference '%s' -> %s", field_name, tool_ref.tool_module_info)
         except Exception:
             logger.exception("Failed to resolve ToolReference '%s'", field_name)
 
     @classmethod
     async def _resolve_list_items(
-        cls, items: list, registry: "RegistryStrategy", communication: "CommunicationStrategy"
+        cls,
+        items: list,
+        registry: "RegistryStrategy",
+        communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve ToolReference instances in a list.
 
@@ -426,16 +426,32 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             items: List of items to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         for item in items:
             if isinstance(item, ToolReference):
-                await cls._resolve_single_tool_reference("list_item", item, registry, communication)
+                await cls._resolve_single_tool_reference(
+                    "list_item",
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
             elif isinstance(item, BaseModel):
-                await cls._resolve_tool_references_recursive(item, registry, communication)
+                await cls._resolve_tool_references_recursive(
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
 
     @classmethod
     async def _resolve_dict_values(
-        cls, mapping: dict, registry: "RegistryStrategy", communication: "CommunicationStrategy"
+        cls,
+        mapping: dict,
+        registry: "RegistryStrategy",
+        communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve ToolReference instances in dict values.
 
@@ -443,63 +459,86 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             mapping: Dict to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         for item in mapping.values():
             if isinstance(item, ToolReference):
-                await cls._resolve_single_tool_reference("dict_value", item, registry, communication)
+                await cls._resolve_single_tool_reference(
+                    "dict_value",
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
             elif isinstance(item, BaseModel):
-                await cls._resolve_tool_references_recursive(item, registry, communication)
+                await cls._resolve_tool_references_recursive(
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
 
     def build_tool_cache(self) -> ToolCache:
-        """Build tool cache from resolved ToolReferences, populating companion fields.
+        """Build tool cache from resolved ToolReferences.
 
         Returns:
             ToolCache with field names as keys and ToolModuleInfo as values.
         """
-        logger.info("Building tool cache")
         cache = ToolCache()
         self._build_tool_cache_recursive(self, cache)
         logger.info("Tool cache built: %d entries", len(cache.entries))
         return cache
 
-    def _build_tool_cache_recursive(self, model_instance: BaseModel, cache: ToolCache) -> None:  # noqa: C901
-        """Recursively build tool cache and populate companion fields.
+    def _build_tool_cache_recursive(self, model_instance: BaseModel, cache: ToolCache) -> None:
+        """Recursively build tool cache from ToolReferences.
 
         Args:
             model_instance: Model instance to process.
             cache: ToolCache to populate.
         """
-        for field_name, field_value in model_instance.__dict__.items():
+        for field_value in model_instance.__dict__.values():
             if field_value is None:
                 continue
             if isinstance(field_value, ToolReference):
-                cache_field_name = f"{field_name}_cache"
-
-                cached_info = getattr(model_instance, cache_field_name, None)
-                module_info = field_value.tool_module_info or cached_info
+                module_info = self.resolved_tools.get(field_value.slug or "") or field_value.tool_module_info
                 if module_info:
-                    if not cached_info:
-                        setattr(model_instance, cache_field_name, module_info)
+                    self.resolved_tools[module_info.setup_id] = module_info
                     cache.add(module_info.module_id, module_info)
-                    logger.debug("Added tool to cache: %s", module_info.module_id)
             elif isinstance(field_value, BaseModel):
                 self._build_tool_cache_recursive(field_value, cache)
             elif isinstance(field_value, list):
-                cache_field_name = f"{field_name}_cache"
-                cached_infos = getattr(model_instance, cache_field_name, None) or []
-                resolved_infos: list[ToolModuleInfo] = []
-
-                for idx, item in enumerate(field_value):
-                    if isinstance(item, ToolReference):
-                        # Use resolved info or fallback to cached
-                        module_info = item.tool_module_info or (cached_infos[idx] if idx < len(cached_infos) else None)
-                        if module_info:
-                            resolved_infos.append(module_info)
-                            cache.add(module_info.module_id, module_info)
-                            logger.debug("Added tool to cache: %s", module_info.module_id)
-                    elif isinstance(item, BaseModel):
-                        self._build_tool_cache_recursive(item, cache)
-
-                # Update companion field with resolved infos
-                if resolved_infos:
-                    setattr(model_instance, cache_field_name, resolved_infos)
+                self._process_list_items(field_value, cache)
+            elif isinstance(field_value, dict):
+                self._process_dict_values(field_value, cache)
+
+    def _process_list_items(self, items: list, cache: ToolCache) -> None:
+        """Process list items for ToolReferences.
+
+        Args:
+            items: List to process.
+            cache: ToolCache to populate.
+        """
+        for item in items:
+            if isinstance(item, ToolReference):
+                module_info = self.resolved_tools.get(item.slug or "") or item.tool_module_info
+                if module_info:
+                    self.resolved_tools[module_info.setup_id] = module_info
+                    cache.add(module_info.module_id, module_info)
+            elif isinstance(item, BaseModel):
+                self._build_tool_cache_recursive(item, cache)
+
+    def _process_dict_values(self, mapping: dict, cache: ToolCache) -> None:
+        """Process dict values for ToolReferences.
+
+        Args:
+            mapping: Dict to process.
+            cache: ToolCache to populate.
+        """
+        for item in mapping.values():
+            if isinstance(item, ToolReference):
+                module_info = self.resolved_tools.get(item.slug or "") or item.tool_module_info
+                if module_info:
+                    self.resolved_tools[module_info.setup_id] = module_info
+                    cache.add(module_info.module_id, module_info)
+            elif isinstance(item, BaseModel):
+                self._build_tool_cache_recursive(item, cache)

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/src/digitalkin/models/module/tool_cache.py

@@ -52,6 +52,7 @@ class ToolModuleInfo(ModuleInfo):
     """Module info for tool modules."""
 
     tools: list[ToolDefinition]
+    setup_id: str
 
 
 class ToolCache(BaseModel):
@@ -59,22 +60,22 @@ class ToolCache(BaseModel):
 
     entries: dict[str, ToolModuleInfo] = Field(default_factory=dict)
 
-    def add(self, setup_tool_name: str, tool_module_info: ToolModuleInfo) -> None:
+    def add(self, setup_id: str, tool_module_info: ToolModuleInfo) -> None:
         """Add a tool to the cache.
 
         Args:
-            setup_tool_name: Field name from SetupModel used as cache key.
+            setup_id: Field name from SetupModel used as cache key.
             tool_module_info: Resolved tool module information.
         """
-        self.entries[setup_tool_name] = tool_module_info
+        self.entries[setup_id] = tool_module_info
         logger.debug(
             "Tool cached",
-            extra={"setup_tool_name": setup_tool_name, "module_id": tool_module_info.module_id},
+            extra={"setup_id": setup_id, "module_id": tool_module_info.module_id},
         )
 
     async def get(
         self,
-        setup_tool_name: str,
+        setup_id: str,
         *,
         registry: RegistryStrategy | None = None,
         communication: "CommunicationStrategy | None" = None,
@@ -82,28 +83,28 @@ class ToolCache(BaseModel):
         """Get a tool from cache, optionally querying registry on miss.
 
         Args:
-            setup_tool_name: Field name to look up.
+            setup_id: Field name to look up.
             registry: Optional registry to query on cache miss.
             communication: Optional communication strategy for schema fetching.
 
         Returns:
             ToolModuleInfo if found, None otherwise.
         """
-        cached = self.entries.get(setup_tool_name)
+        cached = self.entries.get(setup_id)
         if cached:
             return cached
 
         if registry and communication:
             try:
-                setup_info = registry.get_setup(setup_tool_name)
+                setup_info = registry.get_setup(setup_id)
                 if setup_info and setup_info.module_id:
                     info = registry.discover_by_id(setup_info.module_id)
                     if info:
-                        tool_info = await module_info_to_tool_module_info(info, communication)
-                        self.add(setup_tool_name, tool_info)
+                        tool_info = await module_info_to_tool_module_info(info, setup_id, communication)
+                        self.add(setup_id, tool_info)
                         return tool_info
             except Exception:
-                logger.exception("Registry lookup failed", extra={"setup_tool_name": setup_tool_name})
+                logger.exception("Registry lookup failed", extra={"setup_id": setup_id})
 
         return None
 
@@ -122,6 +123,7 @@ class ToolCache(BaseModel):
 
 async def module_info_to_tool_module_info(
     module_info: ModuleInfo,
+    setup_id: str,
     communication: "CommunicationStrategy",
     *,
     llm_format: bool = True,
@@ -133,6 +135,7 @@ async def module_info_to_tool_module_info(
 
     Args:
         module_info: Module info from registry.
+        setup_id: Setup ID from tool configuration.
         communication: Communication strategy for gRPC calls.
         llm_format: Use LLM-friendly schema format.
 
@@ -161,6 +164,7 @@ async def module_info_to_tool_module_info(
         documentation=module_info.documentation,
         status=module_info.status,
         tools=tools,
+        setup_id=setup_id,
     )
 
 

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/src/digitalkin/models/module/tool_reference.py

@@ -115,7 +115,7 @@ class ToolReference(BaseModel):
                 self.config.module_id = setup.module_id
                 info = registry.discover_by_id(self.config.module_id)
                 if info:
-                    tool_module_info = await module_info_to_tool_module_info(info, communication)
+                    tool_module_info = await module_info_to_tool_module_info(info, self.config.setup_id, communication)
                     self._cached_info = tool_module_info
                     return tool_module_info
 
@@ -126,7 +126,9 @@ class ToolReference(BaseModel):
                 organization_id=self.config.organization_id,
             )
             if results:
-                tool_module_info = await module_info_to_tool_module_info(results[0], communication)
+                tool_module_info = await module_info_to_tool_module_info(
+                    results[0], self.config.setup_id, communication
+                )
                 self._cached_info = tool_module_info
                 self.config.module_id = tool_module_info.module_id
                 return tool_module_info

{digitalkin-0.3.2.dev17 → digitalkin-0.3.2.dev19}/src/digitalkin/modules/_base_module.py

@@ -556,21 +556,7 @@ class BaseModule(  # noqa: PLR0904
             await self._resolve_tools(config_setup_data)
             updated_config = await self.run_config_setup(self.context, config_setup_data)
 
-            # Build wrapper: original structure with updated content
-            wrapper = config_setup_data.model_dump()
-            wrapper["content"] = updated_config.model_dump()
-
-            # Debug logging
-            content = wrapper.get("content", {})
-            logger.info(
-                "Config setup wrapper: keys=%s, content_keys=%s, tools_cache=%s",
-                list(wrapper.keys()),
-                list(content.keys()) if isinstance(content, dict) else "N/A",
-                content.get("tools_cache") if isinstance(content, dict) else "N/A",
-                extra=self.context.session.current_ids(),
-            )
-
-            setup_model = await self.create_setup_model(wrapper)
+            setup_model = await self.create_setup_model(updated_config.model_dump())
             await callback(setup_model)
             self._status = ModuleStatus.STOPPING
         except Exception: