diff --git a/opax-mcp/persistence/firestore_store.py b/opax-mcp/persistence/firestore_store.py index 513a7a1..26cbada 100644 --- a/opax-mcp/persistence/firestore_store.py +++ b/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) diff --git a/opax-mcp/persistence/memory_store.py b/opax-mcp/persistence/memory_store.py index 1102df0..6cf5fe6 100644 --- a/opax-mcp/persistence/memory_store.py +++ b/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: diff --git a/opax-mcp/test_phase2a_foundation.py b/opax-mcp/test_phase2a_foundation.py index bdacd60..dac819a 100644 --- a/opax-mcp/test_phase2a_foundation.py +++ b/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.""" diff --git a/opax-mcp/test_phase2b_firestore_store.py b/opax-mcp/test_phase2b_firestore_store.py index 5a132e4..27abf3d 100644 --- a/opax-mcp/test_phase2b_firestore_store.py +++ b/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) + + 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,32 +55,111 @@ 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.""" + 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_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 - convo_collection_mock.document.assert_called_with(convo_id) - doc_mock = convo_collection_mock.document.return_value - doc_mock.collection.assert_called_with("messages") + mock_convo_collection.document.assert_called_with(convo_id) + mock_doc.collection.assert_called_with("messages") # Check that add was called on the subcollection - messages_collection_mock = doc_mock.collection.return_value - messages_collection_mock.add.assert_awaited_once_with(message) + mock_messages_collection.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.""" event = AuditEvent(event_type="test", caller_id="c1", details={}, workspace_id="ws1") await self.store.record_audit_event(event) - + self.mock_firestore_client.collection.assert_called_with("emma_audit_events") collection_mock = self.mock_firestore_client.collection.return_value collection_mock.document.assert_called_with(event.event_id) @@ -76,7 +174,7 @@ class TestPhase2BMemoryScope(unittest.TestCase): """(D) Tests the effective memory scope calculation.""" admin_caller = derive_caller_context("user:chris.christiansen@vauco.no") operator_caller = derive_caller_context("agent:perplexity") - + convo = MagicMock() # Admin requests all, gets all allowed for admin