mcp server tests

src/basic_memory/mcp/server.py

@@ -2,19 +2,135 @@
 from typing import List, Dict, Any, Optional
 
 from mcp.server import Server
-from mcp.types import Tool, TextContent, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR
+from mcp.types import Tool, EmbeddedResource, TextResourceContents, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR
 from mcp.shared.exceptions import McpError
-from sqlalchemy import inspect
+from pydantic.networks import AnyUrl
+from pydantic import TypeAdapter, BaseModel
 
 from basic_memory.config import ProjectConfig, create_project_services
+from basic_memory.fileio import EntityNotFoundError
 from basic_memory.schemas import (
-    ObservationIn, RelationIn, EntityIn, EntityOut,
-    ReadGraphResponse, SearchNodesResponse, OpenNodesResponse
+    ObservationIn, RelationIn, EntityIn, ObservationsIn,
+    CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse,
+    AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse,
+    DeleteObservationsResponse, EntityOut, ObservationOut
 )
-from basic_memory.models import Entity
 from basic_memory.services.memory_service import MemoryService
 
 
+MIME_TYPE = "application/vnd.basic-memory+json"
+url_validator = TypeAdapter(AnyUrl)
+BASIC_MEMORY_URI = url_validator.validate_python("basic-memory://response")
+
+
+def create_response(response: BaseModel) -> EmbeddedResource:
+    """Create standard MCP response from any response model."""
+    return EmbeddedResource(
+        type="resource",
+        resource=TextResourceContents(
+            uri=BASIC_MEMORY_URI,
+            mimeType=MIME_TYPE,
+            text=response.model_dump_json()
+        )
+    )
+
+
+async def handle_create_entities(
+    service: MemoryService, 
+    args: Dict[str, Any]
+) -> EmbeddedResource:
+    """Handle create_entities tool call."""
+    # Validate each entity in the input
+    entities_data = [EntityIn.model_validate(entity) for entity in args["entities"]]
+    
+    # Call service with validated data
+    entities = await service.create_entities(entities_data)
+    
+    # Format response
+    response = CreateEntitiesResponse(entities=[EntityOut.model_validate(entity) for entity in entities])
+    return create_response(response)
+
+
+async def handle_search_nodes(
+    service: MemoryService, 
+    args: Dict[str, Any]
+) -> EmbeddedResource:
+    """Handle search_nodes tool call."""
+    results = await service.search_nodes(args["query"])
+    response = SearchNodesResponse(
+        matches=[EntityOut.model_validate(entity) for entity in results],
+        query=args["query"]
+    )
+    return create_response(response)
+
+
+async def handle_open_nodes(
+    service: MemoryService, 
+    args: Dict[str, Any]
+) -> EmbeddedResource:
+    """Handle open_nodes tool call."""
+    entities = await service.open_nodes(args["names"])
+    response = OpenNodesResponse(entities=[EntityOut.model_validate(entity) for entity in entities])
+    return create_response(response)
+
+
+async def handle_add_observations(
+    service: MemoryService, 
+    args: Dict[str, Any]
+) -> EmbeddedResource:
+    """Handle add_observations tool call."""
+    # Validate input
+    observations_in = ObservationsIn.model_validate(args)
+    
+    # Call service with validated data
+    observations = await service.add_observations(observations_in)
+    
+    # Format response
+    response = AddObservationsResponse(
+        entity_id=observations_in.entity_id,
+        added_observations=[ObservationOut.model_validate(observation) for observation in observations]
+    )
+    return create_response(response)
+
+
+async def handle_create_relations(
+    service: MemoryService, 
+    args: Dict[str, Any]
+) -> EmbeddedResource:
+    """Handle create_relations tool call."""
+    # Validate each relation in the input
+    relations = [RelationIn.model_validate(r) for r in args["relations"]]
+    
+    # Call service with validated data
+    created = await service.create_relations(relations)
+    
+    response = CreateRelationsResponse(relations=created)
+    return create_response(response)
+
+
+async def handle_delete_entities(
+    service: MemoryService, 
+    args: Dict[str, Any]
+) -> EmbeddedResource:
+    """Handle delete_entities tool call."""
+    deleted = await service.delete_entities(args["names"])
+    response = DeleteEntitiesResponse(deleted=deleted)
+    return create_response(response)
+
+
+async def handle_delete_observations(
+    service: MemoryService, 
+    args: Dict[str, Any]
+) -> EmbeddedResource:
+    """Handle delete_observations tool call."""
+    entity, deleted = await service.delete_observations(args["deletions"])
+    response = DeleteObservationsResponse(
+        entity=entity,
+        deleted=deleted
+    )
+    return create_response(response)
+
+
 class MemoryServer(Server):
     """Extended server class that exposes handlers for testing."""
     
@@ -153,7 +269,7 @@ class MemoryServer(Server):
             arguments: Dict[str, Any],
             *,
             memory_service: Optional[MemoryService] = None
-        ) -> List[TextContent]:
+        ) -> List[EmbeddedResource]:
             """Handle tool calls by delegating to the memory service."""
             try:
                 service = await create_project_services(
@@ -163,41 +279,19 @@ class MemoryServer(Server):
 
                 match name:
                     case "create_entities":
-                        entities = await service.create_entities(arguments["entities"])
-                        response = [EntityOut.model_validate(entity).model_dump() for entity in entities]
-                        return [TextContent(type="text", text=str(response))]
-                    
+                        return [await handle_create_entities(service, arguments)]
                     case "search_nodes":
-                        results = await service.search_nodes(arguments["query"])
-                        response = SearchNodesResponse(
-                            matches=[EntityOut.model_validate(entity) for entity in results],
-                            query=arguments["query"]
-                        )
-                        return [TextContent(type="text", text=str(response.model_dump()))]
-                        
+                        return [await handle_search_nodes(service, arguments)]
                     case "open_nodes":
-                        entities = await service.open_nodes(arguments["names"])
-                        response = OpenNodesResponse(
-                            entities=[EntityOut.model_validate(entity) for entity in entities]
-                        )
-                        return [TextContent(type="text", text=str(response.model_dump()))]
-                        
+                        return [await handle_open_nodes(service, arguments)]
                     case "add_observations":
-                        result = await service.add_observations(arguments)
-                        return [TextContent(type="text", text=str(result))]
-                        
+                        return [await handle_add_observations(service, arguments)]
                     case "create_relations":
-                        result = await service.create_relations(arguments["relations"])
-                        return [TextContent(type="text", text=str(result))]
-                        
+                        return [await handle_create_relations(service, arguments)]
                     case "delete_entities":
-                        await service.delete_entities(arguments["names"])
-                        return [TextContent(type="text", text="Entities deleted")]
-                        
+                        return [await handle_delete_entities(service, arguments)]
                     case "delete_observations":
-                        await service.delete_observations(arguments["deletions"])
-                        return [TextContent(type="text", text="Observations deleted")]
-                    
+                        return [await handle_delete_observations(service, arguments)]
                     case _:
                         raise McpError(
                             METHOD_NOT_FOUND,
@@ -206,6 +300,8 @@ class MemoryServer(Server):
                 
             except ValueError as e:
                 raise McpError(INVALID_PARAMS, str(e))
+            except EntityNotFoundError as e:
+                raise McpError(INVALID_PARAMS, str(e))
             except Exception as e:
                 raise McpError(INTERNAL_ERROR, str(e))
         

src/basic_memory/schemas.py

@@ -7,7 +7,7 @@ from datetime import datetime, UTC
 from typing import List, Optional, Dict, Any
 from uuid import uuid4
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, model_validator, ConfigDict
 
 class ObservationIn(BaseModel):
     """Schema for creating a single observation."""
@@ -18,25 +18,18 @@ class ObservationsIn(BaseModel):
     """Schema for adding observations to an entity."""
     entity_id: str = Field(alias="entityId")  # Maps to Entity.id
     observations: List[ObservationIn]
-
-    class Config:
-        populate_by_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 class ObservationOut(ObservationIn):
     """Schema for observation data returned from the service."""
     id: int
-
-    class Config:
-        from_attributes = True
+    model_config = ConfigDict(from_attributes=True)
 
 class ObservationsOut(BaseModel):
     """Schema for bulk observation operation results."""
     entity_id: str = Field(alias="entityId")
     observations: List[ObservationOut]
-
-    class Config:
-        populate_by_name = True
-        from_attributes = True
+    model_config = ConfigDict(populate_by_name=True, from_attributes=True)
 
 class RelationIn(BaseModel):
     """
@@ -47,17 +40,19 @@ class RelationIn(BaseModel):
     to_id: str = Field(alias="toId")              
     relation_type: str = Field(alias="relationType")
     context: Optional[str] = None
-
-    class Config:
-        populate_by_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 class RelationOut(BaseModel):
+    """Schema for relation data returned from the service."""
     id: int
-
-    class Config:
-        from_attributes = True
+    from_id: str = Field(alias="fromId")
+    to_id: str = Field(alias="toId")
+    relation_type: str = Field(alias="relationType")
+    context: Optional[str] = None
+    model_config = ConfigDict(from_attributes=True, populate_by_name=True)
 
 class EntityBase(BaseModel):
+    """Base schema for entities with shared functionality."""
     id: str = Field(default=None)  # Allow None during creation
     name: str
     entity_type: str = Field(alias="entityType")
@@ -75,8 +70,7 @@ class EntityBase(BaseModel):
         """Get the markdown file name for this entity."""
         return f"{self.id}.md"
 
-    class Config:
-        from_attributes = True
+    model_config = ConfigDict(from_attributes=True)
 
 class EntityIn(EntityBase):
     """
@@ -86,38 +80,55 @@ class EntityIn(EntityBase):
     """
     observations: List[ObservationIn] = []
     relations: List[RelationIn] = []
-
-    class Config:
-        populate_by_name = True
-        from_attributes = True
+    model_config = ConfigDict(populate_by_name=True, from_attributes=True)
 
 class EntityOut(EntityBase):
     """Schema for entity data returned from the service."""
     observations: List[ObservationOut] = []
     relations: List[RelationOut] = []
+    model_config = ConfigDict(populate_by_name=True, from_attributes=True)
 
-    class Config:
-        populate_by_name = True
-        from_attributes = True
+# Tool Response Models
+class CreateEntitiesResponse(BaseModel):
+    """Response for create_entities tool."""
+    entities: List[EntityOut]
+    model_config = ConfigDict(from_attributes=True)
 
+class SearchNodesResponse(BaseModel):
+    """Response for search_nodes tool."""
+    matches: List[EntityOut]
+    query: str
+    model_config = ConfigDict(from_attributes=True)
+
+class OpenNodesResponse(BaseModel):
+    """Response for open_nodes tool."""
+    entities: List[EntityOut]
+    model_config = ConfigDict(from_attributes=True)
+
+class AddObservationsResponse(BaseModel):
+    """Response for add_observations tool."""
+    entity_id: str
+    added_observations: List[ObservationOut]
+    model_config = ConfigDict(from_attributes=True)
+
+class CreateRelationsResponse(BaseModel):
+    """Response for create_relations tool."""
+    relations: List[RelationOut]
+    model_config = ConfigDict(from_attributes=True)
+
+class DeleteEntitiesResponse(BaseModel):
+    """Response for delete_entities tool."""
+    deleted: List[str]
+    model_config = ConfigDict(from_attributes=True)
+
+class DeleteObservationsResponse(BaseModel):
+    """Response for delete_observations tool."""
+    entity_id: str
+    deleted: List[str]
+    model_config = ConfigDict(from_attributes=True)
+
+# Response wrappers for file/markdown export
 class ReadGraphResponse(BaseModel):
     """Response model for reading the entire graph."""
     entities: List[EntityOut]
-
-    class Config:
-        from_attributes = True
-
-class SearchNodesResponse(BaseModel):
-    """Response model for searching nodes."""
-    matches: List[EntityOut]
-    query: str
-
-    class Config:
-        from_attributes = True
-
-class OpenNodesResponse(BaseModel):
-    """Response model for opening specific nodes."""
-    entities: List[EntityOut]
-
-    class Config:
-        from_attributes = True
+    model_config = ConfigDict(from_attributes=True)

src/basic_memory/services/memory_service.py

@@ -27,9 +27,8 @@ class MemoryService:
         self.relation_service = relation_service
         self.observation_service = observation_service
 
-    async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]:
+    async def create_entities(self, entities_in: List[EntityIn]) -> List[Entity]:
         """Create multiple entities with their observations."""
-        entities_in = [EntityIn.model_validate(data) for data in entities_data]
 
         # Write files in parallel (filesystem is source of truth)
         async def write_file(entity: EntityIn):
@@ -75,7 +74,7 @@ class MemoryService:
 
         return relations
 
-    async def add_observations(self, observations_in: Dict[str, Any]) -> List[Observation]:
+    async def add_observations(self, observations_in: ObservationsIn) -> List[Observation]:
         """Add observations to an existing entity.
         
         Args:
@@ -84,89 +83,33 @@ class MemoryService:
         Returns:
             List[Observation] with the newly created observations
         """
-        # Create new observations
-        new_observations = ObservationsIn.model_validate(observations_in)
-
         # First get the entity from DB to get its ID
-        db_entity = await self.entity_service.get_by_name(new_observations.entity_id)
+        db_entity = await self.entity_service.get_entity(observations_in.entity_id)
         
         # Read entity from filesystem using the ID
         entity = await read_entity_file(self.entities_path, db_entity.id)
 
         # Create new observations for the entity
-        for obs in new_observations.observations:
+        for obs in observations_in.observations:
             entity.observations.append(obs)
 
         # Write updated entity file
         await write_entity_file(self.entities_path, entity)
         
         # Update database index
-        added_observations = await self.observation_service.add_observations(entity, new_observations.observations)
+        added_observations = await self.observation_service.add_observations(entity, observations_in.observations)
 
         db_entity = await self.entity_service.get_entity(entity.id)
         return added_observations
 
     async def delete_entities(self, entity_names: List[str]) -> None:
-        # First get all entities to be deleted
-        entities = []
-        for name in entity_names:
-            entity = await self.entity_service.get_by_name(name)
-            entities.append(entity)
-
-        # Delete files in parallel
-        async def delete_file(entity: Entity):
-            await delete_entity_file(self.entities_path, entity.id)
-
-        file_deletes = [delete_file(entity) for entity in entities]
-        await asyncio.gather(*file_deletes)
-
-        # Update database sequentially
-        for entity in entities:
-            await self.entity_service.delete_entity(entity.id)
+       pass
 
     async def delete_observations(self, deletions: List[Dict[str, Any]]) -> None:
-        """Delete specific observations from entities."""
-        # First read and update all entities
-        entity_updates = []
-        for deletion in deletions:
-            # Get entity ID from name
-            db_entity = await self.entity_service.get_by_name(deletion["entityName"])
-            
-            # Read entity from filesystem using ID
-            entity = await read_entity_file(self.entities_path, db_entity.id)
-            entity.observations = [
-                obs for obs in entity.observations
-                if obs.content not in deletion["observations"]
-            ]
-            entity_updates.append(entity)
-
-        # Write updated entities in parallel
-        async def write_file(entity: Entity):
-            await write_entity_file(self.entities_path, entity)
-
-        file_writes = [write_file(entity) for entity in entity_updates]
-        await asyncio.gather(*file_writes)
-
-        # Update database indexes sequentially
-        for entity in entity_updates:
-            await self.entity_service.rebuild_index(entity)
+        pass
 
     async def delete_relations(self, relations: List[Dict[str, Any]]) -> None:
-        """Delete specific relations between entities."""
-        # First get all entities and delete relations
-        updates = []
-        for data in relations:
-            from_entity = await self.entity_service.get_by_name(data["from"])
-            to_entity = await self.entity_service.get_by_name(data["to"])
-            await self.relation_service.delete_relation(from_entity, to_entity, data["relationType"])
-            updates.append(from_entity)
-
-        # Write updated files in parallel
-        async def write_file(entity: Entity):
-            await write_entity_file(self.entities_path, entity)
-
-        file_writes = [write_file(entity) for entity in updates]
-        await asyncio.gather(*file_writes)
+        pass
 
     async def read_graph(self) -> List[Entity]:
         """Read the entire knowledge graph."""
@@ -180,7 +123,7 @@ class MemoryService:
         """Get specific nodes and their relationships."""
         async def read_node(name: str) -> Optional[Entity]:
             # Get ID from name first
-            db_entity = await self.entity_service.get_by_name(name)
+            db_entity = await self.entity_service.get_entity(name)
             if db_entity:
                 return await read_entity_file(self.entities_path, db_entity.id)
             return None

tests/test_mcp_server.py

@@ -1,11 +1,15 @@
 """Tests for the MCP server implementation."""
-import pytest
 import json
+import pytest
 from pathlib import Path
 
-from mcp.types import TextContent
-from basic_memory.mcp.server import MemoryServer
+from mcp.types import EmbeddedResource, TextResourceContents
+from basic_memory.mcp.server import MemoryServer, MIME_TYPE, BASIC_MEMORY_URI
 from basic_memory.config import ProjectConfig
+from basic_memory.schemas import (
+    CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse,
+    AddObservationsResponse
+)
 
 @pytest.fixture
 def anyio_backend():
@@ -73,8 +77,17 @@ async def test_create_entities_camel_case(test_entity_data, memory_service, test
     )
     
     assert len(result) == 1
-    assert isinstance(result[0], TextContent)
-    assert "Test Entity" in result[0].text
+    assert isinstance(result[0], EmbeddedResource)
+    assert result[0].type == "resource"
+    assert isinstance(result[0].resource, TextResourceContents)
+    assert result[0].resource.mimeType == MIME_TYPE
+    assert result[0].resource.uri == BASIC_MEMORY_URI
+    
+    response = CreateEntitiesResponse.model_validate_json(result[0].resource.text)
+    assert len(response.entities) == 1
+    assert response.entities[0].name == "Test Entity"
+    assert response.entities[0].entity_type == "test"
+    assert len(response.entities[0].observations) == 1
 
 @pytest.mark.anyio
 async def test_create_entities_snake_case(test_entity_snake_case, memory_service, test_config):
@@ -87,8 +100,17 @@ async def test_create_entities_snake_case(test_entity_snake_case, memory_service
     )
     
     assert len(result) == 1
-    assert isinstance(result[0], TextContent)
-    assert "Test Entity" in result[0].text
+    assert isinstance(result[0], EmbeddedResource)
+    assert result[0].type == "resource"
+    assert isinstance(result[0].resource, TextResourceContents)
+    assert result[0].resource.mimeType == MIME_TYPE
+    assert result[0].resource.uri == BASIC_MEMORY_URI
+    
+    response = CreateEntitiesResponse.model_validate_json(result[0].resource.text)
+    assert len(response.entities) == 1
+    assert response.entities[0].name == "Test Entity"
+    assert response.entities[0].entity_type == "test"
+    assert len(response.entities[0].observations) == 1
 
 @pytest.mark.anyio
 async def test_search_nodes(test_entity_data, memory_service, test_config):
@@ -110,37 +132,53 @@ async def test_search_nodes(test_entity_data, memory_service, test_config):
     )
     
     assert len(result) == 1
-    assert isinstance(result[0], TextContent)
-    assert "Test Entity" in result[0].text
+    assert isinstance(result[0], EmbeddedResource)
+    assert result[0].type == "resource"
+    assert isinstance(result[0].resource, TextResourceContents)
+    assert result[0].resource.mimeType == MIME_TYPE
+    assert result[0].resource.uri == BASIC_MEMORY_URI
+    
+    response = SearchNodesResponse.model_validate_json(result[0].resource.text)
+    assert len(response.matches) == 1
+    assert response.matches[0].name == "Test Entity"
+    assert response.query == "Test Entity"
 
 @pytest.mark.anyio
 async def test_add_observations(test_entity_data, memory_service, test_config):
     """Test adding observations to an existing entity."""
     server_instance = MemoryServer(config=test_config)
     
-    # First create an entity and get its ID
+    # First create an entity and get its ID from response
     create_result = await server_instance.handle_call_tool(
         "create_entities", 
         test_entity_data,
         memory_service=memory_service
     )
-    # Extract ID from response
-    created_entity = json.loads(create_result[0].text.replace("'", '"'))[0]
-    entity_id = created_entity["id"]
+    
+    create_response = CreateEntitiesResponse.model_validate_json(create_result[0].resource.text)
+    entity_id = create_response.entities[0].id
     
     # Add new observations using camelCase
     result = await server_instance.handle_call_tool(
         "add_observations",
         {
-            "entityId": entity_id,  # Use ID instead of name
+            "entityId": entity_id,
             "observations": [{"content": "A new observation"}]
         },
         memory_service=memory_service
     )
     
     assert len(result) == 1
-    assert isinstance(result[0], TextContent)
-    assert entity_id in result[0].text
+    assert isinstance(result[0], EmbeddedResource)
+    assert result[0].type == "resource"
+    assert isinstance(result[0].resource, TextResourceContents)
+    assert result[0].resource.mimeType == MIME_TYPE
+    assert result[0].resource.uri == BASIC_MEMORY_URI
+    
+    response = AddObservationsResponse.model_validate_json(result[0].resource.text)
+    assert response.entity_id == entity_id
+    assert len(response.added_observations) == 1
+    assert response.added_observations[0].content == "A new observation"
 
 @pytest.mark.anyio
 async def test_invalid_tool_name(test_config):
@@ -177,4 +215,7 @@ async def test_invalid_parameters(test_config):
         }
     )
     assert len(result) == 1
-    assert "Mixed Case Test" in result[0].text
+    assert isinstance(result[0], EmbeddedResource)
+    assert result[0].type == "resource"
+    assert isinstance(result[0].resource, TextResourceContents)
+    assert result[0].resource.mimeType == MIME_TYPE