diff --git a/examples/server/memory/kb.go b/examples/server/memory/kb.go index 2277c22b..e28a4909 100644 --- a/examples/server/memory/kb.go +++ b/examples/server/memory/kb.go @@ -482,34 +482,34 @@ func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerReque }, nil } -func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteEntitiesArgs) (*mcp.CallToolResult, struct{}, error) { +func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteEntitiesArgs) (*mcp.CallToolResult, any, error) { var res mcp.CallToolResult err := k.deleteEntities(args.EntityNames) if err != nil { - return nil, struct{}{}, err + return nil, nil, err } res.Content = []mcp.Content{ &mcp.TextContent{Text: "Entities deleted successfully"}, } - return &res, struct{}{}, nil + return &res, nil, nil } -func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteObservationsArgs) (*mcp.CallToolResult, struct{}, error) { +func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteObservationsArgs) (*mcp.CallToolResult, any, error) { var res mcp.CallToolResult err := k.deleteObservations(args.Deletions) if err != nil { - return nil, struct{}{}, err + return nil, nil, err } res.Content = []mcp.Content{ &mcp.TextContent{Text: "Observations deleted successfully"}, } - return &res, struct{}{}, nil + return &res, nil, nil } func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args DeleteRelationsArgs) (*mcp.CallToolResult, struct{}, error) { @@ -554,8 +554,7 @@ func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[* &mcp.TextContent{Text: "Nodes searched successfully"}, } - res.StructuredContent = graph - return &res, KnowledgeGraph{}, nil + return &res, graph, nil } func (k knowledgeBase) OpenNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParams], args OpenNodesArgs) (*mcp.CallToolResult, KnowledgeGraph, error) { diff --git a/examples/server/memory/kb_test.go b/examples/server/memory/kb_test.go index d0cf38c0..5d40ae64 100644 --- a/examples/server/memory/kb_test.go +++ b/examples/server/memory/kb_test.go @@ -427,203 +427,161 @@ func TestFileFormatting(t *testing.T) { } // TestMCPServerIntegration tests the knowledge base through MCP server layer. -// func TestMCPServerIntegration(t *testing.T) { -// for name, newStore := range stores() { -// t.Run(name, func(t *testing.T) { -// s := newStore(t) -// kb := knowledgeBase{s: s} - -// // Create mock server session -// ctx := context.Background() -// serverSession := &mcp.ServerSession{} - -// // Test CreateEntities through MCP -// createEntitiesParams := &mcp.CallToolParamsFor[CreateEntitiesArgs]{ -// Arguments: CreateEntitiesArgs{ -// Entities: []Entity{ -// { -// Name: "TestPerson", -// EntityType: "Person", -// Observations: []string{"Likes testing"}, -// }, -// }, -// }, -// } - -// createResult, err := kb.CreateEntities(ctx, requestFor(serverSession, createEntitiesParams)) -// if err != nil { -// t.Fatalf("MCP CreateEntities failed: %v", err) -// } -// if createResult.IsError { -// t.Fatalf("MCP CreateEntities returned error: %v", createResult.Content) -// } -// if len(createResult.StructuredContent.Entities) != 1 { -// t.Errorf("expected 1 entity created, got %d", len(createResult.StructuredContent.Entities)) -// } - -// // Test ReadGraph through MCP -// readParams := &mcp.CallToolParamsFor[struct{}]{} -// readResult, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams)) -// if err != nil { -// t.Fatalf("MCP ReadGraph failed: %v", err) -// } -// if readResult.IsError { -// t.Fatalf("MCP ReadGraph returned error: %v", readResult.Content) -// } -// if len(readResult.StructuredContent.Entities) != 1 { -// t.Errorf("expected 1 entity in graph, got %d", len(readResult.StructuredContent.Entities)) -// } - -// // Test CreateRelations through MCP -// createRelationsParams := &mcp.CallToolParamsFor[CreateRelationsArgs]{ -// Arguments: CreateRelationsArgs{ -// Relations: []Relation{ -// { -// From: "TestPerson", -// To: "Testing", -// RelationType: "likes", -// }, -// }, -// }, -// } - -// relationsResult, err := kb.CreateRelations(ctx, requestFor(serverSession, createRelationsParams)) -// if err != nil { -// t.Fatalf("MCP CreateRelations failed: %v", err) -// } -// if relationsResult.IsError { -// t.Fatalf("MCP CreateRelations returned error: %v", relationsResult.Content) -// } -// if len(relationsResult.StructuredContent.Relations) != 1 { -// t.Errorf("expected 1 relation created, got %d", len(relationsResult.StructuredContent.Relations)) -// } - -// // Test AddObservations through MCP -// addObsParams := &mcp.CallToolParamsFor[AddObservationsArgs]{ -// Arguments: AddObservationsArgs{ -// Observations: []Observation{ -// { -// EntityName: "TestPerson", -// Contents: []string{"Works remotely", "Drinks coffee"}, -// }, -// }, -// }, -// } - -// obsResult, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams)) -// if err != nil { -// t.Fatalf("MCP AddObservations failed: %v", err) -// } -// if obsResult.IsError { -// t.Fatalf("MCP AddObservations returned error: %v", obsResult.Content) -// } -// if len(obsResult.StructuredContent.Observations) != 1 { -// t.Errorf("expected 1 observation result, got %d", len(obsResult.StructuredContent.Observations)) -// } - -// // Test SearchNodes through MCP -// searchParams := &mcp.CallToolParamsFor[SearchNodesArgs]{ -// Arguments: SearchNodesArgs{ -// Query: "coffee", -// }, -// } - -// searchResult, err := kb.SearchNodes(ctx, requestFor(serverSession, searchParams)) -// if err != nil { -// t.Fatalf("MCP SearchNodes failed: %v", err) -// } -// if searchResult.IsError { -// t.Fatalf("MCP SearchNodes returned error: %v", searchResult.Content) -// } -// if len(searchResult.StructuredContent.Entities) != 1 { -// t.Errorf("expected 1 entity from search, got %d", len(searchResult.StructuredContent.Entities)) -// } - -// // Test OpenNodes through MCP -// openParams := &mcp.CallToolParamsFor[OpenNodesArgs]{ -// Arguments: OpenNodesArgs{ -// Names: []string{"TestPerson"}, -// }, -// } - -// openResult, err := kb.OpenNodes(ctx, requestFor(serverSession, openParams)) -// if err != nil { -// t.Fatalf("MCP OpenNodes failed: %v", err) -// } -// if openResult.IsError { -// t.Fatalf("MCP OpenNodes returned error: %v", openResult.Content) -// } -// if len(openResult.StructuredContent.Entities) != 1 { -// t.Errorf("expected 1 entity from open, got %d", len(openResult.StructuredContent.Entities)) -// } - -// // Test DeleteObservations through MCP -// deleteObsParams := &mcp.CallToolParamsFor[DeleteObservationsArgs]{ -// Arguments: DeleteObservationsArgs{ -// Deletions: []Observation{ -// { -// EntityName: "TestPerson", -// Observations: []string{"Works remotely"}, -// }, -// }, -// }, -// } - -// deleteObsResult, err := kb.DeleteObservations(ctx, requestFor(serverSession, deleteObsParams)) -// if err != nil { -// t.Fatalf("MCP DeleteObservations failed: %v", err) -// } -// if deleteObsResult.IsError { -// t.Fatalf("MCP DeleteObservations returned error: %v", deleteObsResult.Content) -// } - -// // Test DeleteRelations through MCP -// deleteRelParams := &mcp.CallToolParamsFor[DeleteRelationsArgs]{ -// Arguments: DeleteRelationsArgs{ -// Relations: []Relation{ -// { -// From: "TestPerson", -// To: "Testing", -// RelationType: "likes", -// }, -// }, -// }, -// } - -// deleteRelResult, err := kb.DeleteRelations(ctx, requestFor(serverSession, deleteRelParams)) -// if err != nil { -// t.Fatalf("MCP DeleteRelations failed: %v", err) -// } -// if deleteRelResult.IsError { -// t.Fatalf("MCP DeleteRelations returned error: %v", deleteRelResult.Content) -// } - -// // Test DeleteEntities through MCP -// deleteEntParams := &mcp.CallToolParamsFor[DeleteEntitiesArgs]{ -// Arguments: DeleteEntitiesArgs{ -// EntityNames: []string{"TestPerson"}, -// }, -// } - -// deleteEntResult, err := kb.DeleteEntities(ctx, requestFor(serverSession, deleteEntParams)) -// if err != nil { -// t.Fatalf("MCP DeleteEntities failed: %v", err) -// } -// if deleteEntResult.IsError { -// t.Fatalf("MCP DeleteEntities returned error: %v", deleteEntResult.Content) -// } - -// // Verify final state -// finalRead, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams)) -// if err != nil { -// t.Fatalf("Final MCP ReadGraph failed: %v", err) -// } -// if len(finalRead.StructuredContent.Entities) != 0 { -// t.Errorf("expected empty graph after deletion, got %d entities", len(finalRead.StructuredContent.Entities)) -// } -// }) -// } -// } +func TestMCPServerIntegration(t *testing.T) { + for name, newStore := range stores() { + t.Run(name, func(t *testing.T) { + s := newStore(t) + kb := knowledgeBase{s: s} + + // Create mock server session + ctx := context.Background() + + createResult, out, err := kb.CreateEntities(ctx, nil, CreateEntitiesArgs{ + Entities: []Entity{ + { + Name: "TestPerson", + EntityType: "Person", + Observations: []string{"Likes testing"}, + }, + }, + }) + if err != nil { + t.Fatalf("MCP CreateEntities failed: %v", err) + } + if createResult.IsError { + t.Fatalf("MCP CreateEntities returned error: %v", createResult.Content) + } + if len(out.Entities) != 1 { + t.Errorf("expected 1 entity created, got %d", len(out.Entities)) + } + + // Test ReadGraph through MCP + readResult, outg, err := kb.ReadGraph(ctx, nil, nil) + if err != nil { + t.Fatalf("MCP ReadGraph failed: %v", err) + } + if readResult.IsError { + t.Fatalf("MCP ReadGraph returned error: %v", readResult.Content) + } + if len(outg.Entities) != 1 { + t.Errorf("expected 1 entity in graph, got %d", len(outg.Entities)) + } + + relationsResult, outr, err := kb.CreateRelations(ctx, nil, CreateRelationsArgs{ + Relations: []Relation{ + { + From: "TestPerson", + To: "Testing", + RelationType: "likes", + }, + }, + }) + if err != nil { + t.Fatalf("MCP CreateRelations failed: %v", err) + } + if relationsResult.IsError { + t.Fatalf("MCP CreateRelations returned error: %v", relationsResult.Content) + } + if len(outr.Relations) != 1 { + t.Errorf("expected 1 relation created, got %d", len(outr.Relations)) + } + + obsResult, outo, err := kb.AddObservations(ctx, nil, AddObservationsArgs{ + Observations: []Observation{ + { + EntityName: "TestPerson", + Contents: []string{"Works remotely", "Drinks coffee"}, + }, + }, + }) + if err != nil { + t.Fatalf("MCP AddObservations failed: %v", err) + } + if obsResult.IsError { + t.Fatalf("MCP AddObservations returned error: %v", obsResult.Content) + } + if len(outo.Observations) != 1 { + t.Errorf("expected 1 observation result, got %d", len(outo.Observations)) + } + + searchResult, outg, err := kb.SearchNodes(ctx, nil, SearchNodesArgs{ + Query: "coffee", + }) + if err != nil { + t.Fatalf("MCP SearchNodes failed: %v", err) + } + if searchResult.IsError { + t.Fatalf("MCP SearchNodes returned error: %v", searchResult.Content) + } + if len(outg.Entities) != 1 { + t.Errorf("expected 1 entity from search, got %d", len(outg.Entities)) + } + + openResult, outg, err := kb.OpenNodes(ctx, nil, OpenNodesArgs{ + Names: []string{"TestPerson"}, + }) + if err != nil { + t.Fatalf("MCP OpenNodes failed: %v", err) + } + if openResult.IsError { + t.Fatalf("MCP OpenNodes returned error: %v", openResult.Content) + } + if len(outg.Entities) != 1 { + t.Errorf("expected 1 entity from open, got %d", len(outg.Entities)) + } + + deleteObsResult, _, err := kb.DeleteObservations(ctx, nil, DeleteObservationsArgs{ + Deletions: []Observation{ + { + EntityName: "TestPerson", + Observations: []string{"Works remotely"}, + }, + }, + }) + if err != nil { + t.Fatalf("MCP DeleteObservations failed: %v", err) + } + if deleteObsResult.IsError { + t.Fatalf("MCP DeleteObservations returned error: %v", deleteObsResult.Content) + } + + deleteRelResult, _, err := kb.DeleteRelations(ctx, nil, DeleteRelationsArgs{ + Relations: []Relation{ + { + From: "TestPerson", + To: "Testing", + RelationType: "likes", + }, + }, + }) + if err != nil { + t.Fatalf("MCP DeleteRelations failed: %v", err) + } + if deleteRelResult.IsError { + t.Fatalf("MCP DeleteRelations returned error: %v", deleteRelResult.Content) + } + + deleteEntResult, _, err := kb.DeleteEntities(ctx, nil, DeleteEntitiesArgs{ + EntityNames: []string{"TestPerson"}, + }) + if err != nil { + t.Fatalf("MCP DeleteEntities failed: %v", err) + } + if deleteEntResult.IsError { + t.Fatalf("MCP DeleteEntities returned error: %v", deleteEntResult.Content) + } + + // Verify final state + _, outg, err = kb.ReadGraph(ctx, nil, nil) + if err != nil { + t.Fatalf("Final MCP ReadGraph failed: %v", err) + } + if len(outg.Entities) != 0 { + t.Errorf("expected empty graph after deletion, got %d entities", len(outg.Entities)) + } + }) + } +} // TestMCPErrorHandling tests error scenarios through MCP layer. func TestMCPErrorHandling(t *testing.T) { diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 28246fd3..0e9cf455 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -60,7 +60,7 @@ func TestStreamableTransports(t *testing.T) { return nil, nil, nil } AddTool(server, &Tool{Name: "hang"}, hang) - AddTool(server, &Tool{Name: "sample"}, func(ctx context.Context, req *ServerRequest[*CallToolParams], args map[string]any) (*CallToolResult, any, error) { + AddTool(server, &Tool{Name: "sample"}, func(ctx context.Context, req *ServerRequest[*CallToolParams], args any) (*CallToolResult, any, error) { // Test that we can make sampling requests during tool handling. // // Try this on both the request context and a background context, so