fix(opax): restore persistence store compatibility

opax-mcp/persistence/firestore_store.py

@@ -5,6 +5,7 @@ from typing import List, Dict, Any, Optional
 from datetime import datetime, timezone
 
 from google.cloud import firestore
+from google.cloud.firestore_v1.async_transaction import async_transactional
 
 from contracts.common import EmmaConversation, MemoryRecord, AuditEvent
 from .memory_store import MemoryStore
@@ -30,6 +31,46 @@ class FirestoreMemoryStore(MemoryStore):
         await doc_ref.set(convo.model_dump(mode='json'))
         return convo
 
+    async def save_conversation(self, conversation: EmmaConversation) -> None:
+        if not conversation.owner_id:
+            raise ValueError("Conversation must have an owner_id.")
+
+        db = self._get_client()
+        doc_ref = db.collection("emma_conversations").document(
+            conversation.conversation_id
+        )
+        conversation_data = conversation.model_dump(mode="json")
+
+        @async_transactional
+        async def _save(transaction) -> None:
+            snapshot = await doc_ref.get(transaction=transaction)
+
+            if not snapshot.exists:
+                transaction.create(doc_ref, conversation_data)
+                return
+
+            existing_data = snapshot.to_dict() or {}
+
+            if existing_data.get("owner_id") != conversation.owner_id:
+                raise PermissionError(
+                    "Cannot modify a conversation owned by another owner."
+                )
+
+            if existing_data.get("workspace_id") != conversation.workspace_id:
+                raise PermissionError(
+                    "Cannot move a conversation to another workspace."
+                )
+
+            conversation_data["owner_id"] = existing_data["owner_id"]
+            conversation_data["workspace_id"] = existing_data["workspace_id"]
+
+            if existing_data.get("created_at") is not None:
+                conversation_data["created_at"] = existing_data["created_at"]
+
+            transaction.update(doc_ref, conversation_data)
+
+        await _save(db.transaction())
+
     async def get_conversation(self, conversation_id: str, owner_id: str) -> Optional[EmmaConversation]:
         db = self._get_client()
         doc_ref = db.collection("emma_conversations").document(conversation_id)

opax-mcp/persistence/memory_store.py

@@ -13,6 +13,10 @@ class MemoryStore(ABC):
     async def create_conversation(self, owner_id: str, workspace_id: str, created_by: str) -> EmmaConversation:
         pass
 
+    @abstractmethod
+    async def save_conversation(self, conversation: EmmaConversation) -> None:
+        pass
+
     @abstractmethod
     async def get_conversation(self, conversation_id: str, owner_id: str) -> Optional[EmmaConversation]:
         pass
@@ -55,6 +59,19 @@ class InMemoryMemoryStore(MemoryStore):
         self._messages[convo.conversation_id] = []
         return convo
 
+    async def save_conversation(self, conversation: EmmaConversation) -> None:
+        if not conversation.owner_id:
+            raise ValueError("Conversation must have an owner_id.")
+        # Prevent an existing conversation from being claimed by a new owner.
+        existing = self._conversations.get(conversation.conversation_id)
+        if existing and existing.owner_id != conversation.owner_id:
+            raise ValueError("Cannot change the owner of an existing conversation.")
+
+        self._conversations[conversation.conversation_id] = conversation
+        # Ensure the message list is initialized.
+        if conversation.conversation_id not in self._messages:
+            self._messages[conversation.conversation_id] = []
+
     async def get_conversation(self, conversation_id: str, owner_id: str) -> Optional[EmmaConversation]:
         convo = self._conversations.get(conversation_id)
         if convo and convo.owner_id == owner_id:

opax-mcp/test_phase2a_foundation.py

@@ -100,9 +100,9 @@ class TestPhase2AMemoryStore(unittest.IsolatedAsyncioTestCase):
         await self.store.save_conversation(convo1)
 
         # Owner 1 can retrieve
-        self.assertIsNotNone(await self.store.get_conversation("c1", "owner1"))
+        self.assertIsNotNone(await self.store.get_conversation(convo1.conversation_id, "owner1"))
         # Owner 2 cannot retrieve
-        self.assertIsNone(await self.store.get_conversation("c1", "owner2"))
+        self.assertIsNone(await self.store.get_conversation(convo1.conversation_id, "owner2"))
 
     async def test_approval_record_instantiation(self):
         """(C.2) Tests that ApprovalRecord can be created with a real ProposedToolAction."""

opax-mcp/test_phase2b_firestore_store.py

@@ -19,12 +19,31 @@ from policy.caller_context import derive_caller_context
 class TestPhase2BFirestoreStore(unittest.IsolatedAsyncioTestCase):
 
     def setUp(self):
+        # Patch the transactional decorator to simply execute the inner function
+        self.transaction_patcher = patch("persistence.firestore_store.async_transactional", new=lambda f: f)
+        self.transaction_patcher.start()
+
         self.mock_firestore_client = MagicMock()
+
+        # Mocks for save_conversation
+        self.mock_transaction = MagicMock()
+        self.mock_transaction.create = MagicMock()
+        self.mock_transaction.update = MagicMock()
+        self.mock_firestore_client.transaction.return_value = self.mock_transaction
+
+        # Mocks for create_conversation
         self.mock_firestore_client.collection.return_value.document.return_value.set = AsyncMock()
-        self.mock_firestore_client.collection.return_value.add = AsyncMock()
+
+        # Mocks for append_message
+        self.mock_firestore_client.collection.return_value.document.return_value.collection.return_value.add = AsyncMock()
+        self.mock_firestore_client.collection.return_value.document.return_value.update = AsyncMock()
+
+        # Mocks for record_audit_event
         self.mock_firestore_client.collection.return_value.document.return_value.create = AsyncMock()
 
                 self.store = FirestoreMemoryStore(project_id="test-project", client=self.mock_firestore_client)
+    def tearDown(self):
+        self.transaction_patcher.stop()
 
     async def test_firestore_client_is_injected(self):
         """(1) Verifies the Firestore client is used when provided."""
@@ -36,26 +55,105 @@ class TestPhase2BFirestoreStore(unittest.IsolatedAsyncioTestCase):
         
         self.mock_firestore_client.collection.assert_called_with("emma_conversations")
         collection_mock = self.mock_firestore_client.collection.return_value
+
                 collection_mock.document.assert_called_with(convo.conversation_id)
                 document_mock = collection_mock.document.return_value
                 document_mock.set.assert_awaited_once()
 
-    async def test_append_message_uses_subcollection(self):
+
-        """(3) Verifies messages are stored as a subcollection."""
+    async def test_append_message_uses_subcollection_and_updates_timestamp(self):
+        """(3) Verifies messages are stored as a subcollection and the parent is updated."""
         convo_id = "convo-123"
         message = {"role": "user", "content": "hello"}
+
+        # Configure the mock chain for this specific test
+        mock_convo_collection = self.mock_firestore_client.collection.return_value
+        mock_doc = mock_convo_collection.document.return_value
+        mock_messages_collection = mock_doc.collection.return_value
+
         await self.store.append_message(convo_id, message)
 
         # Check path to messages subcollection
         self.mock_firestore_client.collection.assert_called_with("emma_conversations")
-        convo_collection_mock = self.mock_firestore_client.collection.return_value
+        mock_convo_collection.document.assert_called_with(convo_id)
-        convo_collection_mock.document.assert_called_with(convo_id)
+        mock_doc.collection.assert_called_with("messages")
-        doc_mock = convo_collection_mock.document.return_value
-        doc_mock.collection.assert_called_with("messages")
 
         # Check that add was called on the subcollection
-        messages_collection_mock = doc_mock.collection.return_value
+        mock_messages_collection.add.assert_awaited_once_with(message)
-        messages_collection_mock.add.assert_awaited_once_with(message)
+
+        # Check that the parent conversation's timestamp was updated
+        mock_doc.update.assert_awaited_once()
+        # Ensure the argument contains the last_updated_at field.
+        # We don't care about the exact time, just that it's a datetime.
+        update_call_args = mock_doc.update.call_args[0][0]
+        self.assertIn("last_updated_at", update_call_args)
+        self.assertIsInstance(update_call_args["last_updated_at"], datetime)
+
+    async def test_save_conversation_creates_new_if_not_exists(self):
+        """Tests that save_conversation creates a document if none exists."""
+        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
+        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value
+
+        # Mock the snapshot to show the document doesn't exist
+        mock_snapshot = MagicMock()
+        mock_snapshot.exists = False
+        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)
+
+        await self.store.save_conversation(convo)
+
+        # Verify that create was called on the transaction
+        self.mock_transaction.create.assert_called_once_with(mock_doc_ref, convo.model_dump(mode="json"))
+        self.mock_transaction.update.assert_not_called()
+
+    async def test_save_conversation_updates_existing_with_same_owner(self):
+        """Tests that save_conversation updates a document if it exists with the same owner."""
+        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
+        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value
+
+        # Mock the snapshot to show a matching existing document
+        mock_snapshot = MagicMock()
+        mock_snapshot.exists = True
+        mock_snapshot.to_dict.return_value = {"owner_id": "owner1", "workspace_id": "ws1"}
+        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)
+
+        await self.store.save_conversation(convo)
+
+        self.mock_transaction.update.assert_called_once()
+        self.mock_transaction.create.assert_not_called()
+
+    async def test_save_conversation_raises_error_for_different_owner(self):
+        """Tests that save_conversation fails if the owner_id mismatches."""
+        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
+        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value
+
+        # Mock snapshot with a different owner
+        mock_snapshot = MagicMock()
+        mock_snapshot.exists = True
+        mock_snapshot.to_dict.return_value = {"owner_id": "DIFFERENT_OWNER", "workspace_id": "ws1"}
+        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)
+
+        with self.assertRaises(PermissionError):
+            await self.store.save_conversation(convo)
+
+        self.mock_transaction.create.assert_not_called()
+        self.mock_transaction.update.assert_not_called()
+
+    async def test_save_conversation_raises_error_for_different_workspace(self):
+        """Tests that save_conversation fails if the workspace_id mismatches."""
+        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
+        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value
+
+        # Mock snapshot with a different workspace
+        mock_snapshot = MagicMock()
+        mock_snapshot.exists = True
+        mock_snapshot.to_dict.return_value = {"owner_id": "owner1", "workspace_id": "DIFFERENT_WORKSPACE"}
+        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)
+
+        with self.assertRaises(PermissionError):
+            await self.store.save_conversation(convo)
+
+        self.mock_transaction.create.assert_not_called()
+        self.mock_transaction.update.assert_not_called()
 
     async def test_audit_event_is_append_only(self):
         """(E) Verifies audit events use create() to be append-only."""