Fix(tool cache) : clear resolved_tools in BaseModule for live tool-module schema refetching

.bumpversion.toml

@@ -3,7 +3,7 @@
 # https://peps.python.org/pep-0440/
 
 [tool.bumpversion]
-    current_version = "0.4.2"
+    current_version = "0.4.3"
     parse = """(?x)
     (?P<major>0|[1-9]\\d*)\\.
     (?P<minor>0|[1-9]\\d*)\\.

pyproject.toml

@@ -28,14 +28,14 @@
 
     dependencies = [
         "ag-ui-protocol>=0.1.14",
-        "agentic-mesh-protocol==0.2.3",
+        "agentic-mesh-protocol==0.2.4",
         "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.4.2"
+    version = "0.4.3"
 
     [project.optional-dependencies]
         profiling = [

src/digitalkin/__version__.py

@@ -5,4 +5,4 @@
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.4.2"
+    __version__ = "0.4.3"

src/digitalkin/modules/_base_module.py

@@ -649,6 +649,10 @@ async def _resolve_tools(self, config_setup_data: SetupModelT) -> None:
             config_setup_data: Setup data containing tool references.
         """
         logger.debug("Starting tool resolution", extra=self.context.session.current_ids())
+        # New setup version: discard any inherited resolved_tools so the live
+        # tool-module schemas are re-fetched. Mission runs reuse the persisted
+        # resolved_tools (via build_tool_cache in start()) and never reach here.
+        config_setup_data.resolved_tools = {}
         tool_cache = await config_setup_data.build_tool_cache(self.context.registry, self.context.communication)
         self.context.tool_cache = tool_cache
         logger.debug(

src/digitalkin/services/storage/default_storage.py

@@ -62,7 +62,7 @@ def _load_from_file(self) -> dict[str, StorageRecord]:
                     continue
                 data_model = model_cls.model_validate(rd["data"])
                 rec = StorageRecord(
-                    mission_id=rd["mission_id"],
+                    context=rd["context"],
                     collection=rd["collection"],
                     record_id=rd["record_id"],
                     data=data_model,
@@ -93,7 +93,7 @@ def _save_to_file(self) -> None:
                 serial: dict[str, dict] = {}
                 for key, record in self.storage.items():
                     serial[key] = {
-                        "mission_id": record.mission_id,
+                        "context": record.context,
                         "collection": record.collection,
                         "record_id": record.record_id,
                         "data_type": record.data_type.name,
@@ -119,7 +119,7 @@ async def _store(self, record: StorageRecord) -> StorageRecord:
         Raises:
             ValueError: If the record already exists
         """
-        key = f"{record.collection}:{record.record_id}"
+        key = self._key(record.context, record.collection, record.record_id)
         if key in self.storage:
             msg = f"Document {key!r} already exists"
             raise ValueError(msg)
@@ -131,31 +131,36 @@ async def _store(self, record: StorageRecord) -> StorageRecord:
         logger.debug("Created %s", key)
         return record
 
-    async def _read(self, collection: str, record_id: str) -> StorageRecord | None:
-        """Get records from the database.
+    @staticmethod
+    def _key(context: str, collection: str, record_id: str) -> str:
+        return f"{context}|{collection}:{record_id}"
+
+    async def _read(self, collection: str, record_id: str, context: str) -> StorageRecord | None:
+        """Get a record from the database scoped to a specific context.
 
         Args:
             collection: The unique name to retrieve data for
             record_id: The unique ID of the record
+            context: Owner context scoping the lookup.
 
         Returns:
             StorageRecord: The corresponding record
         """
-        key = f"{collection}:{record_id}"
-        return self.storage.get(key)
+        return self.storage.get(self._key(context, collection, record_id))
 
-    async def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
-        """Update records in the database and persist to file.
+    async def _update(self, collection: str, record_id: str, data: BaseModel, context: str) -> StorageRecord | None:
+        """Update a record in the database scoped to a specific context.
 
         Args:
             collection: The unique name to retrieve data for
             record_id: The unique ID of the record
             data: The data to modify
+            context: Owner context scoping the update.
 
         Returns:
             StorageRecord: The modified record
         """
-        key = f"{collection}:{record_id}"
+        key = self._key(context, collection, record_id)
         rec = self.storage.get(key)
         if not rec:
             return None
@@ -165,46 +170,49 @@ async def _update(self, collection: str, record_id: str, data: BaseModel) -> Sto
         logger.debug("Modified %s", key)
         return rec
 
-    async def _remove(self, collection: str, record_id: str) -> bool:
-        """Delete records from the database and update file.
+    async def _remove(self, collection: str, record_id: str, context: str) -> bool:
+        """Delete a record from the database scoped to a specific context.
 
         Args:
             collection: The unique name to retrieve data for
             record_id: The unique ID of the record
+            context: Owner context scoping the deletion.
 
         Returns:
             bool: True if the record was removed, False otherwise
         """
-        key = f"{collection}:{record_id}"
+        key = self._key(context, collection, record_id)
         if key not in self.storage:
             return False
         del self.storage[key]
         self._save_to_file()
         logger.debug("Removed %s", key)
         return True
 
-    async def _list(self, collection: str) -> list[StorageRecord]:
-        """Implements StorageStrategy._list.
+    async def _list(self, collection: str, context: str) -> list[StorageRecord]:
+        """List records in a collection scoped to a specific context.
 
         Args:
             collection: The unique name to retrieve data for
+            context: Owner context scoping the listing.
 
         Returns:
             A list of storage records
         """
-        prefix = f"{collection}:"
+        prefix = f"{context}|{collection}:"
         return [r for k, r in self.storage.items() if k.startswith(prefix)]
 
-    async def _remove_collection(self, collection: str) -> bool:
-        """Implements StorageStrategy._remove_collection.
+    async def _remove_collection(self, collection: str, context: str) -> bool:
+        """Wipe a collection scoped to a specific context.
 
         Args:
             collection: The unique name to retrieve data for
+            context: Owner context scoping the wipe.
 
         Returns:
             bool: True if the collection was removed, False otherwise
         """
-        prefix = f"{collection}:"
+        prefix = f"{context}|{collection}:"
         to_delete = [k for k in self.storage if k.startswith(prefix)]
         for k in to_delete:
             del self.storage[k]

src/digitalkin/services/storage/grpc_storage.py

@@ -34,7 +34,7 @@ def _build_record_from_proto(self, proto: data_pb2.StorageRecord) -> StorageReco
             A fully validated StorageRecord.
         """
         # Direct field access for scalars (avoids full MessageToDict overhead)
-        mission = proto.mission_id
+        ctx = proto.context
         coll = proto.collection
         rid = proto.record_id
         dtype = DataType[data_pb2.DataType.Name(proto.data_type)]
@@ -48,7 +48,7 @@ def _build_record_from_proto(self, proto: data_pb2.StorageRecord) -> StorageReco
 
         validated = self._validate_data(coll, payload)
         return StorageRecord(
-            mission_id=mission,
+            context=ctx,
             collection=coll,
             record_id=rid,
             data=validated,
@@ -75,7 +75,7 @@ async def _store(self, record: StorageRecord) -> StorageRecord:
             data_struct.update(record.data.model_dump())
             req = data_pb2.StoreRecordRequest(
                 data=data_struct,
-                mission_id=record.mission_id,
+                context=record.context,
                 collection=record.collection,
                 record_id=record.record_id,
                 data_type=record.data_type.name,
@@ -90,16 +90,16 @@ async def _store(self, record: StorageRecord) -> StorageRecord:
             )
             raise StorageServiceError(str(e)) from e
 
-    async def _read(self, collection: str, record_id: str) -> StorageRecord | None:
-        """Fetch a single document by collection + record_id.
+    async def _read(self, collection: str, record_id: str, context: str) -> StorageRecord | None:
+        """Fetch a single document scoped to a specific context.
 
         Returns:
             StorageData: The record
         """
-        logger.debug("debug:_read collection=%s id=%s", collection, record_id)
+        logger.debug("debug:_read context=%s collection=%s id=%s", context, collection, record_id)
         try:
             req = data_pb2.ReadRecordRequest(
-                mission_id=self.mission_id,
+                context=context,
                 collection=collection,
                 record_id=record_id,
             )
@@ -114,24 +114,20 @@ async def _update(
         collection: str,
         record_id: str,
         data: BaseModel,
+        context: str,
     ) -> StorageRecord | None:
-        """Overwrite a document via gRPC.
-
-        Args:
-            collection: The unique name for the record type
-            record_id: The unique ID for the record
-            data: The validated data model
+        """Overwrite a document via gRPC scoped to a specific context.
 
         Returns:
-            StorageRecord: The updated record
+            StorageRecord: The updated record, or None on failure.
         """
-        logger.debug("debug:_update collection=%s id=%s", collection, record_id)
+        logger.debug("debug:_update context=%s collection=%s id=%s", context, collection, record_id)
         try:
             struct = Struct()
             struct.update(data.model_dump())
             req = data_pb2.UpdateRecordRequest(
                 data=struct,
-                mission_id=self.mission_id,
+                context=context,
                 collection=collection,
                 record_id=record_id,
             )
@@ -141,20 +137,16 @@ async def _update(
             logger.warning("gRPC UpdateRecord failed for %s:%s", collection, record_id)
             return None
 
-    async def _remove(self, collection: str, record_id: str) -> bool:
-        """Delete a document via gRPC.
-
-        Args:
-            collection: The unique name for the record type
-            record_id: The unique ID for the record
+    async def _remove(self, collection: str, record_id: str, context: str) -> bool:
+        """Delete a document via gRPC scoped to a specific context.
 
         Returns:
-            bool: True if the record was deleted, False otherwise
+            bool: True if the record was deleted, False otherwise.
         """
-        logger.debug("debug:_remove collection=%s id=%s", collection, record_id)
+        logger.debug("debug:_remove context=%s collection=%s id=%s", context, collection, record_id)
         try:
             req = data_pb2.RemoveRecordRequest(
-                mission_id=self.mission_id,
+                context=context,
                 collection=collection,
                 record_id=record_id,
             )
@@ -168,19 +160,16 @@ async def _remove(self, collection: str, record_id: str) -> bool:
             return False
         return True
 
-    async def _list(self, collection: str) -> list[StorageRecord]:
-        """List all documents in a collection via gRPC.
-
-        Args:
-            collection: The unique name for the record type
+    async def _list(self, collection: str, context: str) -> list[StorageRecord]:
+        """List all documents in a collection via gRPC scoped to a specific context.
 
         Returns:
-            list[StorageRecord]: A list of storage records
+            list[StorageRecord]: The records found, or an empty list on failure.
         """
-        logger.debug("debug:_list collection=%s", collection)
+        logger.debug("debug:_list context=%s collection=%s", context, collection)
         try:
             req = data_pb2.ListRecordsRequest(
-                mission_id=self.mission_id,
+                context=context,
                 collection=collection,
             )
             resp = await self.exec_grpc_query("ListRecords", req)
@@ -189,18 +178,15 @@ async def _list(self, collection: str) -> list[StorageRecord]:
             logger.warning("gRPC ListRecords failed for %s", collection)
             return []
 
-    async def _remove_collection(self, collection: str) -> bool:
-        """Delete an entire collection via gRPC.
-
-        Args:
-            collection: The unique name for the record type
+    async def _remove_collection(self, collection: str, context: str) -> bool:
+        """Delete an entire collection via gRPC scoped to a specific context.
 
         Returns:
-            bool: True if the collection was deleted, False otherwise
+            bool: True if the collection was removed, False otherwise.
         """
         try:
             req = data_pb2.RemoveCollectionRequest(
-                mission_id=self.mission_id,
+                context=context,
                 collection=collection,
             )
             await self.exec_grpc_query("RemoveCollection", req)