mcp: memory server example

examples/memory/kb.go

@@ -15,6 +15,34 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
+// Entity represents a knowledge graph node with observations.
+type Entity struct {
+	Name         string   `json:"name"`
+	EntityType   string   `json:"entityType"`
+	Observations []string `json:"observations"`
+}
+
+// Relation represents a directed edge between two entities.
+type Relation struct {
+	From         string `json:"from"`
+	To           string `json:"to"`
+	RelationType string `json:"relationType"`
+}
+
+// Observation contains facts about an entity.
+type Observation struct {
+	EntityName string   `json:"entityName"`
+	Contents   []string `json:"contents"`
+
+	Observations []string `json:"observations,omitempty"` // Used for deletion operations
+}
+
+// KnowledgeGraph represents the complete graph structure.
+type KnowledgeGraph struct {
+	Entities  []Entity   `json:"entities"`
+	Relations []Relation `json:"relations"`
+}
+
 // store provides persistence interface for knowledge base data.
 type store interface {
 	Read() ([]byte, error)
@@ -154,7 +182,7 @@ func (k knowledgeBase) saveGraph(graph KnowledgeGraph) error {
 }
 
 // createEntities adds new entities to the graph, skipping duplicates by name.
-// Returns the new entities that were actually added.
+// It returns the new entities that were actually added.
 func (k knowledgeBase) createEntities(entities []Entity) ([]Entity, error) {
 	graph, err := k.loadGraph()
 	if err != nil {
@@ -177,7 +205,7 @@ func (k knowledgeBase) createEntities(entities []Entity) ([]Entity, error) {
 }
 
 // createRelations adds new relations to the graph, skipping exact duplicates.
-// Returns the new relations that were actually added.
+// It returns the new relations that were actually added.
 func (k knowledgeBase) createRelations(relations []Relation) ([]Relation, error) {
 	graph, err := k.loadGraph()
 	if err != nil {
@@ -205,7 +233,7 @@ func (k knowledgeBase) createRelations(relations []Relation) ([]Relation, error)
 }
 
 // addObservations appends new observations to existing entities.
-// Returns the new observations that were actually added.
+// It returns the new observations that were actually added.
 func (k knowledgeBase) addObservations(observations []Observation) ([]Observation, error) {
 	graph, err := k.loadGraph()
 	if err != nil {
@@ -408,11 +436,7 @@ func (k knowledgeBase) CreateEntities(ctx context.Context, ss *mcp.ServerSession
 
 	entities, err := k.createEntities(params.Arguments.Entities)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -431,11 +455,7 @@ func (k knowledgeBase) CreateRelations(ctx context.Context, ss *mcp.ServerSessio
 
 	relations, err := k.createRelations(params.Arguments.Relations)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -454,11 +474,7 @@ func (k knowledgeBase) AddObservations(ctx context.Context, ss *mcp.ServerSessio
 
 	observations, err := k.addObservations(params.Arguments.Observations)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -477,11 +493,7 @@ func (k knowledgeBase) DeleteEntities(ctx context.Context, ss *mcp.ServerSession
 
 	err := k.deleteEntities(params.Arguments.EntityNames)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -496,11 +508,7 @@ func (k knowledgeBase) DeleteObservations(ctx context.Context, ss *mcp.ServerSes
 
 	err := k.deleteObservations(params.Arguments.Deletions)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -515,11 +523,7 @@ func (k knowledgeBase) DeleteRelations(ctx context.Context, ss *mcp.ServerSessio
 
 	err := k.deleteRelations(params.Arguments.Relations)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -534,11 +538,7 @@ func (k knowledgeBase) ReadGraph(ctx context.Context, ss *mcp.ServerSession, par
 
 	graph, err := k.loadGraph()
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -554,11 +554,7 @@ func (k knowledgeBase) SearchNodes(ctx context.Context, ss *mcp.ServerSession, p
 
 	graph, err := k.searchNodes(params.Arguments.Query)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{
@@ -574,11 +570,7 @@ func (k knowledgeBase) OpenNodes(ctx context.Context, ss *mcp.ServerSession, par
 
 	graph, err := k.openNodes(params.Arguments.Names)
 	if err != nil {
-		res.IsError = true
-		res.Content = []mcp.Content{
-			&mcp.TextContent{Text: err.Error()},
-		}
-		return &res, nil
+		return nil, err
 	}
 
 	res.Content = []mcp.Content{

examples/memory/kb_test.go

@@ -11,13 +11,15 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"slices"
+	"strings"
 	"testing"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
-// getStoreFactories provides test factories for both storage implementations.
-func getStoreFactories() map[string]func(t *testing.T) store {
+// stores provides test factories for both storage implementations.
+func stores() map[string]func(t *testing.T) store {
 	return map[string]func(t *testing.T) store{
 		"file": func(t *testing.T) store {
 			tempDir, err := os.MkdirTemp("", "kb-test-file-*")
@@ -35,11 +37,9 @@ func getStoreFactories() map[string]func(t *testing.T) store {
 
 // TestKnowledgeBaseOperations verifies CRUD operations work correctly.
 func TestKnowledgeBaseOperations(t *testing.T) {
-	factories := getStoreFactories()
-
-	for name, factory := range factories {
+	for name, newStore := range stores() {
 		t.Run(name, func(t *testing.T) {
-			s := factory(t)
+			s := newStore(t)
 			kb := knowledgeBase{s: s}
 
 			// Verify empty graph loads correctly
@@ -147,19 +147,16 @@ func TestKnowledgeBaseOperations(t *testing.T) {
 
 			// Confirm observation removal
 			graph, _ = kb.loadGraph()
-			aliceFound := false
-			for _, e := range graph.Entities {
-				if e.Name == "Alice" {
-					aliceFound = true
-					for _, obs := range e.Observations {
-						if obs == "Works as developer" {
-							t.Errorf("observation 'Works as developer' should have been deleted")
-						}
-					}
-				}
-			}
-			if !aliceFound {
+			aliceIndex := slices.IndexFunc(graph.Entities, func(e Entity) bool {
+				return e.Name == "Alice"
+			})
+			if aliceIndex == -1 {
 				t.Errorf("entity 'Alice' not found after deleting observation")
+			} else {
+				alice := graph.Entities[aliceIndex]
+				if slices.Contains(alice.Observations, "Works as developer") {
+					t.Errorf("observation 'Works as developer' should have been deleted")
+				}
 			}
 
 			// Remove relations
@@ -191,11 +188,9 @@ func TestKnowledgeBaseOperations(t *testing.T) {
 
 // TestSaveAndLoadGraph ensures data persists correctly across save/load cycles.
 func TestSaveAndLoadGraph(t *testing.T) {
-	factories := getStoreFactories()
-
-	for name, factory := range factories {
+	for name, newStore := range stores() {
 		t.Run(name, func(t *testing.T) {
-			s := factory(t)
+			s := newStore(t)
 			kb := knowledgeBase{s: s}
 
 			// Setup test data
@@ -251,11 +246,9 @@ func TestSaveAndLoadGraph(t *testing.T) {
 
 // TestDuplicateEntitiesAndRelations verifies duplicate prevention logic.
 func TestDuplicateEntitiesAndRelations(t *testing.T) {
-	factories := getStoreFactories()
-
-	for name, factory := range factories {
+	for name, newStore := range stores() {
 		t.Run(name, func(t *testing.T) {
-			s := factory(t)
+			s := newStore(t)
 			kb := knowledgeBase{s: s}
 
 			// Setup initial state
@@ -355,10 +348,9 @@ func TestErrorHandling(t *testing.T) {
 		}
 	})
 
-	factories := getStoreFactories()
-	for name, factory := range factories {
+	for name, newStore := range stores() {
 		t.Run(fmt.Sprintf("AddObservationToNonExistentEntity_%s", name), func(t *testing.T) {
-			s := factory(t)
+			s := newStore(t)
 			kb := knowledgeBase{s: s}
 
 			// Setup valid entity for comparison
@@ -385,11 +377,9 @@ func TestErrorHandling(t *testing.T) {
 
 // TestFileFormatting verifies the JSON storage format structure.
 func TestFileFormatting(t *testing.T) {
-	factories := getStoreFactories()
-
-	for name, factory := range factories {
+	for name, newStore := range stores() {
 		t.Run(name, func(t *testing.T) {
-			s := factory(t)
+			s := newStore(t)
 			kb := knowledgeBase{s: s}
 
 			// Setup test entity
@@ -438,11 +428,9 @@ func TestFileFormatting(t *testing.T) {
 
 // TestMCPServerIntegration tests the knowledge base through MCP server layer.
 func TestMCPServerIntegration(t *testing.T) {
-	factories := getStoreFactories()
-
-	for name, factory := range factories {
+	for name, newStore := range stores() {
 		t.Run(name, func(t *testing.T) {
-			s := factory(t)
+			s := newStore(t)
 			kb := knowledgeBase{s: s}
 
 			// Create mock server session
@@ -639,11 +627,9 @@ func TestMCPServerIntegration(t *testing.T) {
 
 // TestMCPErrorHandling tests error scenarios through MCP layer.
 func TestMCPErrorHandling(t *testing.T) {
-	factories := getStoreFactories()
-
-	for name, factory := range factories {
+	for name, newStore := range stores() {
 		t.Run(name, func(t *testing.T) {
-			s := factory(t)
+			s := newStore(t)
 			kb := knowledgeBase{s: s}
 
 			ctx := context.Background()
@@ -661,15 +647,15 @@ func TestMCPErrorHandling(t *testing.T) {
 				},
 			}
 
-			obsResult, err := kb.AddObservations(ctx, serverSession, addObsParams)
-			if err != nil {
-				t.Fatalf("MCP AddObservations call failed: %v", err)
-			}
-			if !obsResult.IsError {
+			_, err := kb.AddObservations(ctx, serverSession, addObsParams)
+			if err == nil {
 				t.Errorf("expected MCP AddObservations to return error for non-existent entity")
-			}
-			if len(obsResult.Content) == 0 {
-				t.Errorf("expected error content in MCP response")
+			} else {
+				// Verify the error message contains expected text
+				expectedErrorMsg := "entity with name NonExistentEntity not found"
+				if !strings.Contains(err.Error(), expectedErrorMsg) {
+					t.Errorf("expected error message to contain '%s', got: %v", expectedErrorMsg, err)
+				}
 			}
 		})
 	}

examples/memory/main.go

@@ -24,34 +24,6 @@ type HiArgs struct {
 	Name string `json:"name"`
 }
 
-// Entity represents a knowledge graph node with observations.
-type Entity struct {
-	Name         string   `json:"name"`
-	EntityType   string   `json:"entityType"`
-	Observations []string `json:"observations"`
-}
-
-// Relation represents a directed edge between two entities.
-type Relation struct {
-	From         string `json:"from"`
-	To           string `json:"to"`
-	RelationType string `json:"relationType"`
-}
-
-// Observation contains facts about an entity.
-type Observation struct {
-	EntityName string   `json:"entityName"`
-	Contents   []string `json:"contents"`
-
-	Observations []string `json:"observations,omitempty"` // Used for deletion operations
-}
-
-// KnowledgeGraph represents the complete graph structure.
-type KnowledgeGraph struct {
-	Entities  []Entity   `json:"entities"`
-	Relations []Relation `json:"relations"`
-}
-
 // CreateEntitiesArgs defines the create entities tool parameters.
 type CreateEntitiesArgs struct {
 	Entities []Entity `json:"entities" mcp:"entities to create"`