mcp: introduce Requests (#267)

Handlers and client/server methods take a single request argument
combining the session and params. This slightly simplifies the signature
for handlers, since there is no longer an often-ignored session
argument.

But more importantly, it opens the door to adding more information in
requests, such as auth info and HTTP request headers.

For #243.
This commit is contained in:
Jonathan Amsterdam 2025-08-12 17:49:50 -04:00 • committed by GitHub
commit 6e03217c83
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 517 additions and 380 deletions

View file

@ -84,7 +84,7 @@ func (fs *fileStore) Read() ([]byte, error) {
// Write saves data to file with 0600 permissions.
func (fs *fileStore) Write(data []byte) error {
if err := os.WriteFile(fs.path, data, 0600); err != nil {
if err := os.WriteFile(fs.path, data, 0o600); err != nil {
return fmt.Errorf("failed to write file %s: %w", fs.path, err)
}
return nil
@ -431,10 +431,10 @@ func (k knowledgeBase) openNodes(names []string) (KnowledgeGraph, error) {
}, nil
}
func (k knowledgeBase) CreateEntities(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[CreateEntitiesArgs]) (*mcp.CallToolResultFor[CreateEntitiesResult], error) {
func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateEntitiesArgs]]) (*mcp.CallToolResultFor[CreateEntitiesResult], error) {
var res mcp.CallToolResultFor[CreateEntitiesResult]
entities, err := k.createEntities(params.Arguments.Entities)
entities, err := k.createEntities(req.Params.Arguments.Entities)
if err != nil {
return nil, err
}
@ -450,10 +450,10 @@ func (k knowledgeBase) CreateEntities(ctx context.Context, ss *mcp.ServerSession
return &res, nil
}
func (k knowledgeBase) CreateRelations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[CreateRelationsArgs]) (*mcp.CallToolResultFor[CreateRelationsResult], error) {
func (k knowledgeBase) CreateRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateRelationsArgs]]) (*mcp.CallToolResultFor[CreateRelationsResult], error) {
var res mcp.CallToolResultFor[CreateRelationsResult]
relations, err := k.createRelations(params.Arguments.Relations)
relations, err := k.createRelations(req.Params.Arguments.Relations)
if err != nil {
return nil, err
}
@ -469,10 +469,10 @@ func (k knowledgeBase) CreateRelations(ctx context.Context, ss *mcp.ServerSessio
return &res, nil
}
func (k knowledgeBase) AddObservations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[AddObservationsArgs]) (*mcp.CallToolResultFor[AddObservationsResult], error) {
func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddObservationsArgs]]) (*mcp.CallToolResultFor[AddObservationsResult], error) {
var res mcp.CallToolResultFor[AddObservationsResult]
observations, err := k.addObservations(params.Arguments.Observations)
observations, err := k.addObservations(req.Params.Arguments.Observations)
if err != nil {
return nil, err
}
@ -488,10 +488,10 @@ func (k knowledgeBase) AddObservations(ctx context.Context, ss *mcp.ServerSessio
return &res, nil
}
func (k knowledgeBase) DeleteEntities(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[DeleteEntitiesArgs]) (*mcp.CallToolResultFor[struct{}], error) {
func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteEntitiesArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
var res mcp.CallToolResultFor[struct{}]
err := k.deleteEntities(params.Arguments.EntityNames)
err := k.deleteEntities(req.Params.Arguments.EntityNames)
if err != nil {
return nil, err
}
@ -503,10 +503,10 @@ func (k knowledgeBase) DeleteEntities(ctx context.Context, ss *mcp.ServerSession
return &res, nil
}
func (k knowledgeBase) DeleteObservations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[DeleteObservationsArgs]) (*mcp.CallToolResultFor[struct{}], error) {
func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteObservationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
var res mcp.CallToolResultFor[struct{}]
err := k.deleteObservations(params.Arguments.Deletions)
err := k.deleteObservations(req.Params.Arguments.Deletions)
if err != nil {
return nil, err
}
@ -518,10 +518,10 @@ func (k knowledgeBase) DeleteObservations(ctx context.Context, ss *mcp.ServerSes
return &res, nil
}
func (k knowledgeBase) DeleteRelations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[DeleteRelationsArgs]) (*mcp.CallToolResultFor[struct{}], error) {
func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteRelationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
var res mcp.CallToolResultFor[struct{}]
err := k.deleteRelations(params.Arguments.Relations)
err := k.deleteRelations(req.Params.Arguments.Relations)
if err != nil {
return nil, err
}
@ -533,7 +533,7 @@ func (k knowledgeBase) DeleteRelations(ctx context.Context, ss *mcp.ServerSessio
return &res, nil
}
func (k knowledgeBase) ReadGraph(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[struct{}]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
func (k knowledgeBase) ReadGraph(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[struct{}]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
var res mcp.CallToolResultFor[KnowledgeGraph]
graph, err := k.loadGraph()
@ -549,10 +549,10 @@ func (k knowledgeBase) ReadGraph(ctx context.Context, ss *mcp.ServerSession, par
return &res, nil
}
func (k knowledgeBase) SearchNodes(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[SearchNodesArgs]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SearchNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
var res mcp.CallToolResultFor[KnowledgeGraph]
graph, err := k.searchNodes(params.Arguments.Query)
graph, err := k.searchNodes(req.Params.Arguments.Query)
if err != nil {
return nil, err
}
@ -565,10 +565,10 @@ func (k knowledgeBase) SearchNodes(ctx context.Context, ss *mcp.ServerSession, p
return &res, nil
}
func (k knowledgeBase) OpenNodes(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[OpenNodesArgs]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
func (k knowledgeBase) OpenNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[OpenNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
var res mcp.CallToolResultFor[KnowledgeGraph]
graph, err := k.openNodes(params.Arguments.Names)
graph, err := k.openNodes(req.Params.Arguments.Names)
if err != nil {
return nil, err
}

View file

@ -230,7 +230,7 @@ func TestSaveAndLoadGraph(t *testing.T) {
// Test malformed data handling
if fs, ok := s.(*fileStore); ok {
err := os.WriteFile(fs.path, []byte("invalid json"), 0600)
err := os.WriteFile(fs.path, []byte("invalid json"), 0o600)
if err != nil {
t.Fatalf("failed to write invalid json: %v", err)
}
@ -450,7 +450,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
createResult, err := kb.CreateEntities(ctx, serverSession, createEntitiesParams)
createResult, err := kb.CreateEntities(ctx, requestFor(serverSession, createEntitiesParams))
if err != nil {
t.Fatalf("MCP CreateEntities failed: %v", err)
}
@ -463,7 +463,7 @@ func TestMCPServerIntegration(t *testing.T) {
// Test ReadGraph through MCP
readParams := &mcp.CallToolParamsFor[struct{}]{}
readResult, err := kb.ReadGraph(ctx, serverSession, readParams)
readResult, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams))
if err != nil {
t.Fatalf("MCP ReadGraph failed: %v", err)
}
@ -487,7 +487,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
relationsResult, err := kb.CreateRelations(ctx, serverSession, createRelationsParams)
relationsResult, err := kb.CreateRelations(ctx, requestFor(serverSession, createRelationsParams))
if err != nil {
t.Fatalf("MCP CreateRelations failed: %v", err)
}
@ -510,7 +510,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
obsResult, err := kb.AddObservations(ctx, serverSession, addObsParams)
obsResult, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams))
if err != nil {
t.Fatalf("MCP AddObservations failed: %v", err)
}
@ -528,7 +528,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
searchResult, err := kb.SearchNodes(ctx, serverSession, searchParams)
searchResult, err := kb.SearchNodes(ctx, requestFor(serverSession, searchParams))
if err != nil {
t.Fatalf("MCP SearchNodes failed: %v", err)
}
@ -546,7 +546,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
openResult, err := kb.OpenNodes(ctx, serverSession, openParams)
openResult, err := kb.OpenNodes(ctx, requestFor(serverSession, openParams))
if err != nil {
t.Fatalf("MCP OpenNodes failed: %v", err)
}
@ -569,7 +569,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
deleteObsResult, err := kb.DeleteObservations(ctx, serverSession, deleteObsParams)
deleteObsResult, err := kb.DeleteObservations(ctx, requestFor(serverSession, deleteObsParams))
if err != nil {
t.Fatalf("MCP DeleteObservations failed: %v", err)
}
@ -590,7 +590,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
deleteRelResult, err := kb.DeleteRelations(ctx, serverSession, deleteRelParams)
deleteRelResult, err := kb.DeleteRelations(ctx, requestFor(serverSession, deleteRelParams))
if err != nil {
t.Fatalf("MCP DeleteRelations failed: %v", err)
}
@ -605,7 +605,7 @@ func TestMCPServerIntegration(t *testing.T) {
},
}
deleteEntResult, err := kb.DeleteEntities(ctx, serverSession, deleteEntParams)
deleteEntResult, err := kb.DeleteEntities(ctx, requestFor(serverSession, deleteEntParams))
if err != nil {
t.Fatalf("MCP DeleteEntities failed: %v", err)
}
@ -614,7 +614,7 @@ func TestMCPServerIntegration(t *testing.T) {
}
// Verify final state
finalRead, err := kb.ReadGraph(ctx, serverSession, readParams)
finalRead, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams))
if err != nil {
t.Fatalf("Final MCP ReadGraph failed: %v", err)
}
@ -647,7 +647,7 @@ func TestMCPErrorHandling(t *testing.T) {
},
}
_, err := kb.AddObservations(ctx, serverSession, addObsParams)
_, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams))
if err == nil {
t.Errorf("expected MCP AddObservations to return error for non-existent entity")
} else {
@ -678,7 +678,7 @@ func TestMCPResponseFormat(t *testing.T) {
},
}
result, err := kb.CreateEntities(ctx, serverSession, createParams)
result, err := kb.CreateEntities(ctx, requestFor(serverSession, createParams))
if err != nil {
t.Fatalf("CreateEntities failed: %v", err)
}
@ -701,3 +701,7 @@ func TestMCPResponseFormat(t *testing.T) {
t.Errorf("expected Content[0] to be TextContent")
}
}
func requestFor[P mcp.Params](ss *mcp.ServerSession, p P) *mcp.ServerRequest[P] {
return &mcp.ServerRequest[P]{Session: ss, Params: p}
}