@@ -3,7 +3,6 @@ package ghmcp
33import (
44 "context"
55 "fmt"
6- "io"
76 "log"
87 "net/http"
98 "net/url"
@@ -14,12 +13,10 @@ import (
1413
1514 "github.com/github/github-mcp-server/pkg/errors"
1615 "github.com/github/github-mcp-server/pkg/github"
17- mcplog "github.com/github/github-mcp-server/pkg/log"
1816 "github.com/github/github-mcp-server/pkg/raw"
1917 "github.com/github/github-mcp-server/pkg/translations"
2018 gogithub "github.com/google/go-github/v72/github"
21- "github.com/mark3labs/mcp-go/mcp"
22- "github.com/mark3labs/mcp-go/server"
19+ "github.com/modelcontextprotocol/go-sdk/mcp"
2320 "github.com/shurcooL/githubv4"
2421 "github.com/sirupsen/logrus"
2522)
@@ -49,7 +46,7 @@ type MCPServerConfig struct {
4946 Translator translations.TranslationHelperFunc
5047}
5148
52- func NewMCPServer (cfg MCPServerConfig ) (* server. MCPServer , error ) {
49+ func NewMCPServer (cfg MCPServerConfig ) (* mcp. Server , error ) {
5350 apiHost , err := parseAPIHost (cfg .Host )
5451 if err != nil {
5552 return nil , fmt .Errorf ("failed to parse API host: %w" , err )
@@ -72,35 +69,14 @@ func NewMCPServer(cfg MCPServerConfig) (*server.MCPServer, error) {
7269 } // We're going to wrap the Transport later in beforeInit
7370 gqlClient := githubv4 .NewEnterpriseClient (apiHost .graphqlURL .String (), gqlHTTPClient )
7471
75- // When a client send an initialize request, update the user agent to include the client info.
76- beforeInit := func (_ context.Context , _ any , message * mcp.InitializeRequest ) {
77- userAgent := fmt .Sprintf (
78- "github-mcp-server/%s (%s/%s)" ,
79- cfg .Version ,
80- message .Params .ClientInfo .Name ,
81- message .Params .ClientInfo .Version ,
82- )
72+ ghServer := mcp .NewServer (& mcp.Implementation {
73+ Name : "gh-mcp-server" ,
74+ Title : "GitHub MCP Server" ,
75+ Version : cfg .Version ,
76+ }, & mcp.ServerOptions {})
8377
84- restClient .UserAgent = userAgent
85-
86- gqlHTTPClient .Transport = & userAgentTransport {
87- transport : gqlHTTPClient .Transport ,
88- agent : userAgent ,
89- }
90- }
91-
92- hooks := & server.Hooks {
93- OnBeforeInitialize : []server.OnBeforeInitializeFunc {beforeInit },
94- OnBeforeAny : []server.BeforeAnyHookFunc {
95- func (ctx context.Context , _ any , _ mcp.MCPMethod , _ any ) {
96- // Ensure the context is cleared of any previous errors
97- // as context isn't propagated through middleware
98- errors .ContextWithGitHubErrors (ctx )
99- },
100- },
101- }
102-
103- ghServer := github .NewServer (cfg .Version , server .WithHooks (hooks ))
78+ ghServer .AddReceivingMiddleware (InitializeUserAgentMiddleware [* mcp.ServerSession ](cfg , restClient , gqlHTTPClient ))
79+ ghServer .AddReceivingMiddleware (ApiErrorMiddleware [* mcp.ServerSession ]())
10480
10581 enabledToolsets := cfg .EnabledToolsets
10682 if cfg .DynamicToolsets {
@@ -140,14 +116,56 @@ func NewMCPServer(cfg MCPServerConfig) (*server.MCPServer, error) {
140116 // Register all mcp functionality with the server
141117 tsg .RegisterAll (ghServer )
142118
143- if cfg .DynamicToolsets {
144- dynamic := github .InitDynamicToolset (ghServer , tsg , cfg .Translator )
145- dynamic .RegisterTools (ghServer )
119+ return ghServer , nil
120+ }
121+
122+ func InitializeUserAgentMiddleware [S mcp.Session ](cfg MCPServerConfig , restClient * gogithub.Client , gqlHTTPClient * http.Client ) mcp.Middleware [S ] {
123+ return func (next mcp.MethodHandler [S ]) mcp.MethodHandler [S ] {
124+ return func (ctx context.Context , session S , method string , params mcp.Params ) (mcp.Result , error ) {
125+ // If this is an initialize request, we need to set the user agent
126+ if method == "initialize" {
127+ intializeParams , ok := params .(* mcp.InitializeParams )
128+ if ! ok {
129+ return nil , fmt .Errorf ("expected params to be of type *mcp.InitializeParams, got %T" , params )
130+ }
131+
132+ userAgent := fmt .Sprintf (
133+ "github-mcp-server/%s (%s/%s)" ,
134+ cfg .Version ,
135+ intializeParams .ClientInfo .Name ,
136+ intializeParams .ClientInfo .Version ,
137+ )
138+
139+ restClient .UserAgent = userAgent
140+
141+ gqlHTTPClient .Transport = & userAgentTransport {
142+ transport : gqlHTTPClient .Transport ,
143+ agent : userAgent ,
144+ }
145+ }
146+ return next (ctx , session , method , params )
147+ }
146148 }
149+ }
147150
148- return ghServer , nil
151+ func ApiErrorMiddleware [S mcp.Session ]() mcp.Middleware [S ] {
152+ return func (next mcp.MethodHandler [S ]) mcp.MethodHandler [S ] {
153+ return func (ctx context.Context , session S , method string , params mcp.Params ) (mcp.Result , error ) {
154+ // Ensure the context is cleared of any previous errors
155+ // as context isn't propagated through middleware
156+ ctx = errors .ContextWithGitHubErrors (ctx )
157+
158+ // Call the next handler
159+ return next (ctx , session , method , params )
160+ }
161+ }
149162}
150163
164+ // ghServer := github.NewServer(cfg.Version, server.WithHooks(hooks))
165+
166+ // return ghServer, nil
167+ // }
168+
151169type StdioServerConfig struct {
152170 // Version of the server
153171 Version string
@@ -201,8 +219,6 @@ func RunStdioServer(cfg StdioServerConfig) error {
201219 return fmt .Errorf ("failed to create MCP server: %w" , err )
202220 }
203221
204- stdioServer := server .NewStdioServer (ghServer )
205-
206222 logrusLogger := logrus .New ()
207223 if cfg .LogFilePath != "" {
208224 file , err := os .OpenFile (cfg .LogFilePath , os .O_CREATE | os .O_WRONLY | os .O_APPEND , 0600 )
@@ -214,7 +230,6 @@ func RunStdioServer(cfg StdioServerConfig) error {
214230 logrusLogger .SetOutput (file )
215231 }
216232 stdLogger := log .New (logrusLogger .Writer (), "stdioserver" , 0 )
217- stdioServer .SetErrorLogger (stdLogger )
218233
219234 if cfg .ExportTranslations {
220235 // Once server is initialized, all translations are loaded
@@ -224,15 +239,18 @@ func RunStdioServer(cfg StdioServerConfig) error {
224239 // Start listening for messages
225240 errC := make (chan error , 1 )
226241 go func () {
227- in , out := io .Reader (os .Stdin ), io .Writer (os .Stdout )
242+ stdioTransport := mcp .NewStdioTransport ()
243+ loggingTransport := mcp .NewLoggingTransport (stdioTransport , stdLogger .Writer ())
228244
229- if cfg .EnableCommandLogging {
230- loggedIO := mcplog .NewIOLogger (in , out , logrusLogger )
231- in , out = loggedIO , loggedIO
232- }
245+ // in, out := io.Reader(os.Stdin), io.Writer(os.Stdout)
246+
247+ // if cfg.EnableCommandLogging {
248+ // loggedIO := mcplog.NewIOLogger(in, out, logrusLogger)
249+ // in, out = loggedIO, loggedIO
250+ // }
233251 // enable GitHub errors in the context
234252 ctx := errors .ContextWithGitHubErrors (ctx )
235- errC <- stdioServer . Listen (ctx , in , out )
253+ errC <- ghServer . Run (ctx , loggingTransport )
236254 }()
237255
238256 // Output github-mcp-server string
0 commit comments