diff --git a/go.mod b/go.mod index 358a271a7a..b11455139d 100644 --- a/go.mod +++ b/go.mod @@ -10,7 +10,7 @@ require ( github.com/josephburnett/jd/v2 v2.5.0 github.com/lithammer/fuzzysearch v1.1.8 github.com/microcosm-cc/bluemonday v1.0.27 - github.com/modelcontextprotocol/go-sdk v1.7.0-pre.1 + github.com/modelcontextprotocol/go-sdk v1.7.0-pre.2 github.com/muesli/cache2go v0.0.0-20221011235721-518229cd8021 github.com/shurcooL/githubv4 v0.0.0-20240727222349-48295856cce7 github.com/shurcooL/graphql v0.0.0-20230722043721-ed46e5a46466 diff --git a/go.sum b/go.sum index f3f23b549a..64e5b47839 100644 --- a/go.sum +++ b/go.sum @@ -39,8 +39,8 @@ github.com/lithammer/fuzzysearch v1.1.8 h1:/HIuJnjHuXS8bKaiTMeeDlW2/AyIWk2brx1V8 github.com/lithammer/fuzzysearch v1.1.8/go.mod h1:IdqeyBClc3FFqSzYq/MXESsS4S0FsZ5ajtkr5xPLts4= github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= -github.com/modelcontextprotocol/go-sdk v1.7.0-pre.1 h1:GlMIJyMHFX76bBSQuBCLXZ7pB9cGh4VBS6O5wGd0tgI= -github.com/modelcontextprotocol/go-sdk v1.7.0-pre.1/go.mod h1:dL7u98E/zjJTGzEq+j30jQ8K2k1mb6LeAH4inEcSGts= +github.com/modelcontextprotocol/go-sdk v1.7.0-pre.2 h1:3JwUps1pdSpXYndBMGO9SMca6CSkP9AKnOaKAkSSGHc= +github.com/modelcontextprotocol/go-sdk v1.7.0-pre.2/go.mod h1:dL7u98E/zjJTGzEq+j30jQ8K2k1mb6LeAH4inEcSGts= github.com/muesli/cache2go v0.0.0-20221011235721-518229cd8021 h1:31Y+Yu373ymebRdJN1cWLLooHH8xAr0MhKTEJGV/87g= github.com/muesli/cache2go v0.0.0-20221011235721-518229cd8021/go.mod h1:WERUkUryfUWlrHnFSO/BEUZ+7Ns8aZy7iVOGewxKzcc= github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= diff --git a/internal/ghmcp/oauth.go b/internal/ghmcp/oauth.go index abc6d3d11c..35e48f5bbc 100644 --- a/internal/ghmcp/oauth.go +++ b/internal/ghmcp/oauth.go @@ -3,10 +3,13 @@ package ghmcp import ( "context" "crypto/rand" + "errors" "fmt" "log/slog" + "strings" "github.com/github/github-mcp-server/internal/oauth" + "github.com/github/github-mcp-server/pkg/inventory" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -91,41 +94,187 @@ func (p *sessionPrompter) PromptForm(ctx context.Context, prompt oauth.Prompt) e type oauthAuthenticator interface { HasToken() bool Authenticate(ctx context.Context, prompter oauth.Prompter) (*oauth.Outcome, error) + AwaitToken(ctx context.Context, flowID string) (*oauth.Outcome, error) + Cancel(flowID string) bool } -// createOAuthMiddleware returns receiving middleware that authorizes the session -// lazily, on the first tool call. Authorization is deferred until here (rather -// than at startup) because the prompts depend on an initialized session whose -// elicitation capabilities are known. +// oauthElicitIDPrefix identifies authorization responses in the multi-round-trip +// InputResponses map. The suffix is the manager's per-flow ID, which prevents a +// delayed response from an older prompt from affecting a newer flow. +const oauthElicitIDPrefix = "github_authorization:" + +// protocolVersionNoServerElicitation is the first MCP protocol version that +// forbids server-initiated JSON-RPC requests (SEP-2322): from this version on +// the server may not send elicitation/create while serving a request and must +// instead return an InputRequests map from the tool call (multi round-trip +// requests). It mirrors the go-sdk's internal constant of the same value, which +// the SDK does not export. +const protocolVersionNoServerElicitation = "2026-07-28" + +// serverMayInitiateElicitation reports whether the server is permitted to send +// elicitation requests to the client itself, which the spec allows only before +// protocol version 2026-07-28. A nil or un-negotiated session (only reached in +// unit tests; a real tools/call is always initialized) is treated as legacy. +func serverMayInitiateElicitation(ss *mcp.ServerSession) bool { + if ss == nil { + return true + } + params := ss.InitializeParams() + return params == nil || params.ProtocolVersion < protocolVersionNoServerElicitation +} + +// createOAuthToolMiddleware returns tool-handler middleware that authorizes the +// session lazily, on the first tool call. It runs inside the SDK's +// Server.callTool handler so results returned here still receive SDK +// finalization, including resultType: "input_required" for multi-round-trip +// responses. // // When a token is already available the call proceeds untouched. Otherwise the -// flow runs: secure channels (browser, URL elicitation) block until the token -// arrives and then the call proceeds; the last-resort channel returns the -// instruction to the user as a tool result and asks them to retry. -func createOAuthMiddleware(mgr oauthAuthenticator, logger *slog.Logger) func(next mcp.MethodHandler) mcp.MethodHandler { - return func(next mcp.MethodHandler) mcp.MethodHandler { - return func(ctx context.Context, method string, request mcp.Request) (mcp.Result, error) { - if method != "tools/call" || mgr.HasToken() { - return next(ctx, method, request) +// authorization flow runs, presenting its prompt over whichever channel the +// negotiated protocol allows: on protocol versions before 2026-07-28 the server +// elicits directly; from 2026-07-28 on, where server-initiated requests are +// forbidden (SEP-2322), it uses multi-round-trip elicitation returned from the +// tool call. Either way the last-resort channel returns the instruction as a +// tool result and asks the user to retry. +func createOAuthToolMiddleware(mgr oauthAuthenticator, logger *slog.Logger) inventory.ToolHandlerMiddleware { + return func(next mcp.ToolHandler) mcp.ToolHandler { + return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if !serverMayInitiateElicitation(req.Session) { + if flowID, response, ok := authorizationElicitResponse(req.Params.InputResponses); ok { + return resumeMultiRoundTripAuthorization(ctx, mgr, next, req, flowID, response, logger) + } } - callReq, ok := request.(*mcp.CallToolRequest) - if !ok { - return next(ctx, method, request) + if mgr.HasToken() { + return next(ctx, req) } - - outcome, err := mgr.Authenticate(ctx, &sessionPrompter{session: callReq.Session}) - if err != nil { - return nil, fmt.Errorf("github authorization failed: %w", err) + if serverMayInitiateElicitation(req.Session) { + return authorizeViaServerElicitation(ctx, mgr, next, req, logger) } - if outcome != nil && outcome.UserAction != nil { - logger.Info("surfacing github authorization instructions to user") - return &mcp.CallToolResult{ - Content: []mcp.Content{&mcp.TextContent{Text: outcome.UserAction.Message}}, - }, nil - } - return next(ctx, method, request) + return startMultiRoundTripAuthorization(ctx, mgr, next, req, logger) + } + } +} + +// authorizeViaServerElicitation drives authorization on legacy protocol versions +// (before 2026-07-28), where the server may present the prompt itself via +// server-initiated elicitation. It blocks until the token arrives, then proceeds. +func authorizeViaServerElicitation(ctx context.Context, mgr oauthAuthenticator, next mcp.ToolHandler, req *mcp.CallToolRequest, logger *slog.Logger) (*mcp.CallToolResult, error) { + outcome, err := mgr.Authenticate(ctx, &sessionPrompter{session: req.Session}) + if err != nil { + return nil, fmt.Errorf("github authorization failed: %w", err) + } + if outcome != nil && outcome.UserAction != nil { + logger.Info("surfacing github authorization instructions to user") + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: outcome.UserAction.Message}}, + }, nil + } + return next(ctx, req) +} + +// authorizationElicitResponse finds the authorization response and extracts the +// flow ID encoded in its input-request key. +func authorizationElicitResponse(responses mcp.InputResponseMap) (string, *mcp.ElicitResult, bool) { + for id, response := range responses { + flowID, ok := strings.CutPrefix(id, oauthElicitIDPrefix) + if !ok || flowID == "" { + continue + } + result, _ := response.(*mcp.ElicitResult) + return flowID, result, true + } + return "", nil, false +} + +// startMultiRoundTripAuthorization starts authorization on protocol version +// 2026-07-28 or later. Server-initiated requests are forbidden there (SEP-2322), +// so the prompt is returned as an elicitation input request for the client to +// fulfill and retry. +func startMultiRoundTripAuthorization(ctx context.Context, mgr oauthAuthenticator, next mcp.ToolHandler, req *mcp.CallToolRequest, logger *slog.Logger) (*mcp.CallToolResult, error) { + outcome, err := mgr.Authenticate(ctx, nil) + if err != nil { + return nil, fmt.Errorf("github authorization failed: %w", err) + } + if outcome == nil || outcome.UserAction == nil { + // Already authorized (e.g. the server opened a browser and the flow + // completed); proceed. + return next(ctx, req) + } + + elicit := authorizationElicitParams(outcome.UserAction, &sessionPrompter{session: req.Session}) + if elicit == nil || outcome.FlowID == "" { + // The client cannot present an elicitation (no capability, or no URL to + // show), or the flow cannot be correlated; fall back to returning the + // instructions as a tool result. + logger.Info("surfacing github authorization instructions to user") + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: outcome.UserAction.Message}}, + }, nil + } + logger.Info("requesting github authorization via elicitation") + return &mcp.CallToolResult{ + InputRequests: mcp.InputRequestMap{oauthElicitIDPrefix + outcome.FlowID: elicit}, + }, nil +} + +// resumeMultiRoundTripAuthorization handles the client's retry after it +// fulfilled the authorization elicitation. +func resumeMultiRoundTripAuthorization(ctx context.Context, mgr oauthAuthenticator, next mcp.ToolHandler, req *mcp.CallToolRequest, flowID string, response *mcp.ElicitResult, logger *slog.Logger) (*mcp.CallToolResult, error) { + if response == nil || response.Action != "accept" { + if !mgr.Cancel(flowID) { + return expiredAuthorizationResult(), nil } + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "GitHub authorization was declined. Retry when you're ready to authorize."}}, + }, nil + } + + outcome, err := mgr.AwaitToken(ctx, flowID) + if errors.Is(err, oauth.ErrStaleAuthorizationFlow) { + return expiredAuthorizationResult(), nil + } + if err != nil { + return nil, fmt.Errorf("github authorization failed: %w", err) + } + if outcome != nil && outcome.UserAction != nil { + // The user acknowledged the prompt but has not finished authorizing; + // surface the instructions so they can complete it and retry. + logger.Info("surfacing github authorization instructions to user") + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: outcome.UserAction.Message}}, + }, nil + } + return next(ctx, req) +} + +func expiredAuthorizationResult() *mcp.CallToolResult { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "This GitHub authorization prompt has expired. Retry the request to authorize again."}}, + } +} + +// authorizationElicitParams builds the elicitation that presents the +// authorization instructions to the user. It mirrors sessionPrompter's channel +// selection: URL-mode when the client supports it, otherwise form-mode. It +// returns nil when the client advertised no elicitation capability or there is +// no authorization URL to show, so the caller falls back to a tool-result +// message. +func authorizationElicitParams(ua *oauth.UserAction, p *sessionPrompter) *mcp.ElicitParams { + if ua.URL == "" { + return nil + } + message := "Authorize the GitHub MCP Server to continue." + if ua.UserCode != "" { + message = fmt.Sprintf("Enter code %s to authorize the GitHub MCP Server.", ua.UserCode) + } + switch { + case p.CanPromptURL(): + return &mcp.ElicitParams{Mode: "url", Message: message, URL: ua.URL, ElicitationID: rand.Text()} + case p.CanPromptForm(): + return &mcp.ElicitParams{Mode: "form", Message: ua.Message} + default: + return nil } } diff --git a/internal/ghmcp/oauth_test.go b/internal/ghmcp/oauth_test.go index 732d080e40..ca9b6177ac 100644 --- a/internal/ghmcp/oauth_test.go +++ b/internal/ghmcp/oauth_test.go @@ -13,6 +13,7 @@ import ( "github.com/github/github-mcp-server/pkg/github" "github.com/github/github-mcp-server/pkg/http/headers" "github.com/github/github-mcp-server/pkg/utils" + "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -140,57 +141,49 @@ func TestSessionPrompterCapabilities(t *testing.T) { } } -func TestSessionPrompterPromptActions(t *testing.T) { +// TestSessionPrompterModernProtocolUnavailable verifies that on protocol version +// 2026-07-28 and later — the default negotiated by current clients — the server +// may not initiate elicitation (SEP-2322), so PromptURL and PromptForm report +// the prompt as undeliverable. This is what routes authorization to the +// multi-round-trip path instead (see authorizeViaMultiRoundTrip). +func TestSessionPrompterModernProtocolUnavailable(t *testing.T) { t.Parallel() - tests := []struct { - name string - action string - wantDecline bool - }{ - {name: "accept", action: "accept", wantDecline: false}, - {name: "decline", action: "decline", wantDecline: true}, - {name: "cancel", action: "cancel", wantDecline: true}, - } - caps := &mcp.ClientCapabilities{Elicitation: &mcp.ElicitationCapabilities{ URL: &mcp.URLElicitationCapabilities{}, Form: &mcp.FormElicitationCapabilities{}, }} - for _, tc := range tests { - // URL and form modes share the accept/decline mapping; cover both. - for _, mode := range []string{"url", "form"} { - t.Run(tc.name+"/"+mode, func(t *testing.T) { - t.Parallel() - - handler := func(_ context.Context, _ *mcp.ElicitRequest) (*mcp.ElicitResult, error) { - return &mcp.ElicitResult{Action: tc.action}, nil - } + // The handler should never be reached: the SDK blocks the server-initiated + // request before it leaves the server. + handler := func(_ context.Context, _ *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "accept"}, nil + } - got := runProbe(t, caps, handler, func(ctx context.Context, p *sessionPrompter) string { - var err error - if mode == "url" { - err = p.PromptURL(ctx, oauth.Prompt{Message: "msg", URL: "https://example.com/auth"}) - } else { - err = p.PromptForm(ctx, oauth.Prompt{Message: "msg"}) - } - if err == nil { - return "ok" - } - if err == oauth.ErrPromptDeclined { - return "declined" - } - return "error: " + err.Error() - }) + for _, mode := range []string{"url", "form"} { + t.Run(mode, func(t *testing.T) { + t.Parallel() - if tc.wantDecline { - assert.Equal(t, "declined", got) + got := runProbe(t, caps, handler, func(ctx context.Context, p *sessionPrompter) string { + var err error + if mode == "url" { + err = p.PromptURL(ctx, oauth.Prompt{Message: "msg", URL: "https://example.com/auth"}) } else { - assert.Equal(t, "ok", got) + err = p.PromptForm(ctx, oauth.Prompt{Message: "msg"}) + } + switch { + case err == nil: + return "ok" + case errors.Is(err, oauth.ErrPromptUnavailable): + return "unavailable" + default: + return "error: " + err.Error() } }) - } + + assert.Equal(t, "unavailable", got, + "server-initiated elicitation must be reported undeliverable on protocol 2026-07-28+") + }) } } @@ -247,6 +240,18 @@ type fakeAuthenticator struct { err error authCalls int lastPrompter oauth.Prompter + + // awaitOutcome/awaitErr are returned by AwaitToken; tokenAfterAwait flips + // HasToken to true once AwaitToken is called, simulating a flow that + // acquires the token while the user acts on the elicitation. + awaitOutcome *oauth.Outcome + awaitErr error + tokenAfterAwait bool + awaitCalls int + cancelCalls int + cancelResult bool + lastAwaitFlowID string + lastCancelFlowID string } func (f *fakeAuthenticator) HasToken() bool { return f.hasToken } @@ -257,34 +262,38 @@ func (f *fakeAuthenticator) Authenticate(_ context.Context, prompter oauth.Promp return f.outcome, f.err } -func TestCreateOAuthMiddleware(t *testing.T) { +func (f *fakeAuthenticator) AwaitToken(_ context.Context, flowID string) (*oauth.Outcome, error) { + f.awaitCalls++ + f.lastAwaitFlowID = flowID + if f.tokenAfterAwait { + f.hasToken = true + } + return f.awaitOutcome, f.awaitErr +} + +func (f *fakeAuthenticator) Cancel(flowID string) bool { + f.cancelCalls++ + f.lastCancelFlowID = flowID + return f.cancelResult +} + +func TestCreateOAuthToolMiddleware(t *testing.T) { t.Parallel() const nextText = "handler-ran" - newNext := func(called *bool) mcp.MethodHandler { - return func(_ context.Context, _ string, _ mcp.Request) (mcp.Result, error) { + newNext := func(called *bool) mcp.ToolHandler { + return func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) { *called = true return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: nextText}}}, nil } } - t.Run("non tool call passes through without authenticating", func(t *testing.T) { - t.Parallel() - fake := &fakeAuthenticator{hasToken: false} - var called bool - mw := createOAuthMiddleware(fake, discardLogger()) - _, err := mw(newNext(&called))(context.Background(), "initialize", &mcp.InitializeRequest{}) - require.NoError(t, err) - assert.True(t, called, "next should run") - assert.Zero(t, fake.authCalls, "authentication must not run for non tool calls") - }) - t.Run("existing token short circuits authentication", func(t *testing.T) { t.Parallel() fake := &fakeAuthenticator{hasToken: true} var called bool - mw := createOAuthMiddleware(fake, discardLogger()) - _, err := mw(newNext(&called))(context.Background(), "tools/call", &mcp.CallToolRequest{}) + mw := createOAuthToolMiddleware(fake, discardLogger()) + _, err := mw(newNext(&called))(context.Background(), &mcp.CallToolRequest{}) require.NoError(t, err) assert.True(t, called, "next should run") assert.Zero(t, fake.authCalls, "authentication must be skipped when a token already exists") @@ -294,15 +303,13 @@ func TestCreateOAuthMiddleware(t *testing.T) { t.Parallel() fake := &fakeAuthenticator{hasToken: false, outcome: nil, err: nil} var called bool - mw := createOAuthMiddleware(fake, discardLogger()) - res, err := mw(newNext(&called))(context.Background(), "tools/call", &mcp.CallToolRequest{}) + mw := createOAuthToolMiddleware(fake, discardLogger()) + res, err := mw(newNext(&called))(context.Background(), &mcp.CallToolRequest{}) require.NoError(t, err) assert.Equal(t, 1, fake.authCalls) assert.True(t, called, "next should run once authorized") - callRes, ok := res.(*mcp.CallToolResult) - require.True(t, ok) - require.Len(t, callRes.Content, 1) - assert.Equal(t, nextText, callRes.Content[0].(*mcp.TextContent).Text) + require.Len(t, res.Content, 1) + assert.Equal(t, nextText, res.Content[0].(*mcp.TextContent).Text) }) t.Run("pending user action is surfaced as a tool result", func(t *testing.T) { @@ -310,28 +317,227 @@ func TestCreateOAuthMiddleware(t *testing.T) { const message = "Open https://example.com/auth to authorize, then retry." fake := &fakeAuthenticator{hasToken: false, outcome: &oauth.Outcome{UserAction: &oauth.UserAction{Message: message}}} var called bool - mw := createOAuthMiddleware(fake, discardLogger()) - res, err := mw(newNext(&called))(context.Background(), "tools/call", &mcp.CallToolRequest{}) + mw := createOAuthToolMiddleware(fake, discardLogger()) + res, err := mw(newNext(&called))(context.Background(), &mcp.CallToolRequest{}) require.NoError(t, err) assert.False(t, called, "next must not run while the user still needs to authorize") - callRes, ok := res.(*mcp.CallToolResult) - require.True(t, ok) - require.Len(t, callRes.Content, 1) - assert.Equal(t, message, callRes.Content[0].(*mcp.TextContent).Text) + require.Len(t, res.Content, 1) + assert.Equal(t, message, res.Content[0].(*mcp.TextContent).Text) }) t.Run("authentication error is returned", func(t *testing.T) { t.Parallel() fake := &fakeAuthenticator{hasToken: false, err: assert.AnError} var called bool - mw := createOAuthMiddleware(fake, discardLogger()) - _, err := mw(newNext(&called))(context.Background(), "tools/call", &mcp.CallToolRequest{}) + mw := createOAuthToolMiddleware(fake, discardLogger()) + _, err := mw(newNext(&called))(context.Background(), &mcp.CallToolRequest{}) require.Error(t, err) assert.ErrorIs(t, err, assert.AnError) assert.False(t, called, "next must not run when authentication fails") }) } +// runOAuthMiddlewareCall stands up an in-memory client/server pair with the +// OAuth middleware installed ahead of a probe tool, then calls the tool from a +// default (protocol 2026-07-28) client — driving the multi-round-trip +// authorization path. It returns the final tool-result text and whether the +// probe tool ultimately ran. +func runOAuthMiddlewareCall( + t *testing.T, + fake *fakeAuthenticator, + clientCaps *mcp.ClientCapabilities, + elicitationHandler func(context.Context, *mcp.ElicitRequest) (*mcp.ElicitResult, error), +) (string, bool) { + t.Helper() + + server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "v0.0.1"}, nil) + var toolRan bool + handler := createOAuthToolMiddleware(fake, discardLogger())(func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + toolRan = true + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "tool-ran"}}}, nil + }) + server.AddTool(&mcp.Tool{ + Name: probeToolName, + InputSchema: &jsonschema.Schema{Type: "object"}, + }, handler) + + st, ct := mcp.NewInMemoryTransports() + + ss, err := server.Connect(context.Background(), st, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = ss.Close() }) + + client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "v0.0.1"}, &mcp.ClientOptions{ + Capabilities: clientCaps, + ElicitationHandler: elicitationHandler, + }) + cs, err := client.Connect(context.Background(), ct, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = cs.Close() }) + + res, err := cs.CallTool(context.Background(), &mcp.CallToolParams{Name: probeToolName}) + require.NoError(t, err) + require.Len(t, res.Content, 1) + text, ok := res.Content[0].(*mcp.TextContent) + require.True(t, ok, "tool result should be text content") + return text.Text, toolRan +} + +// TestOAuthMiddlewareMultiRoundTrip exercises the protocol-2026-07-28 path, where +// server-initiated elicitation is forbidden and authorization must be presented +// as a multi-round-trip input request that the client fulfills and retries. +func TestOAuthMiddlewareMultiRoundTrip(t *testing.T) { + t.Parallel() + + urlCaps := &mcp.ClientCapabilities{Elicitation: &mcp.ElicitationCapabilities{URL: &mcp.URLElicitationCapabilities{}}} + + t.Run("accepted elicitation authorizes and proceeds", func(t *testing.T) { + t.Parallel() + fake := &fakeAuthenticator{ + outcome: &oauth.Outcome{UserAction: &oauth.UserAction{URL: "https://example.com/auth", Message: "manual"}, FlowID: "flow-1"}, + tokenAfterAwait: true, + } + var elicited int + accept := func(_ context.Context, _ *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + elicited++ + return &mcp.ElicitResult{Action: "accept"}, nil + } + + text, toolRan := runOAuthMiddlewareCall(t, fake, urlCaps, accept) + + assert.Equal(t, "tool-ran", text, "the tool should run once authorization completes") + assert.True(t, toolRan) + assert.Equal(t, 1, elicited, "the client should be asked to authorize exactly once") + assert.Equal(t, 1, fake.awaitCalls, "the middleware should await the token on retry") + assert.Equal(t, "flow-1", fake.lastAwaitFlowID) + assert.Zero(t, fake.cancelCalls) + assert.Nil(t, fake.lastPrompter, "the manager must not be given a prompter on this protocol") + }) + + t.Run("declined elicitation cancels and does not run the tool", func(t *testing.T) { + t.Parallel() + fake := &fakeAuthenticator{ + outcome: &oauth.Outcome{UserAction: &oauth.UserAction{URL: "https://example.com/auth", Message: "manual"}, FlowID: "flow-1"}, + cancelResult: true, + } + decline := func(_ context.Context, _ *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "decline"}, nil + } + + text, toolRan := runOAuthMiddlewareCall(t, fake, urlCaps, decline) + + assert.False(t, toolRan, "the tool must not run when authorization is declined") + assert.Contains(t, text, "declined") + assert.Equal(t, 1, fake.cancelCalls, "a decline should cancel the in-flight flow") + assert.Equal(t, "flow-1", fake.lastCancelFlowID) + assert.Zero(t, fake.awaitCalls) + }) + + t.Run("stale decline does not cancel the current flow", func(t *testing.T) { + t.Parallel() + fake := &fakeAuthenticator{ + outcome: &oauth.Outcome{UserAction: &oauth.UserAction{URL: "https://example.com/auth", Message: "manual"}, FlowID: "old-flow"}, + } + decline := func(_ context.Context, _ *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + return &mcp.ElicitResult{Action: "decline"}, nil + } + + text, toolRan := runOAuthMiddlewareCall(t, fake, urlCaps, decline) + + assert.False(t, toolRan) + assert.Contains(t, text, "expired") + assert.Equal(t, "old-flow", fake.lastCancelFlowID) + assert.Zero(t, fake.awaitCalls) + }) + + t.Run("form-only client receives actionable instructions", func(t *testing.T) { + t.Parallel() + const ( + authURL = "https://example.com/auth" + message = "Open https://example.com/auth and enter code ABCD-1234." + ) + fake := &fakeAuthenticator{ + outcome: &oauth.Outcome{ + UserAction: &oauth.UserAction{URL: authURL, UserCode: "ABCD-1234", Message: message}, + FlowID: "flow-1", + }, + tokenAfterAwait: true, + } + var elicited *mcp.ElicitParams + accept := func(_ context.Context, req *mcp.ElicitRequest) (*mcp.ElicitResult, error) { + elicited = req.Params + return &mcp.ElicitResult{Action: "accept"}, nil + } + formCaps := &mcp.ClientCapabilities{Elicitation: &mcp.ElicitationCapabilities{Form: &mcp.FormElicitationCapabilities{}}} + + text, toolRan := runOAuthMiddlewareCall(t, fake, formCaps, accept) + + assert.True(t, toolRan) + assert.Equal(t, "tool-ran", text) + require.NotNil(t, elicited) + assert.Equal(t, "form", elicited.Mode) + assert.Contains(t, elicited.Message, authURL) + assert.Contains(t, elicited.Message, "ABCD-1234") + }) + + t.Run("no elicitation capability falls back to a tool-result message", func(t *testing.T) { + t.Parallel() + const message = "Open https://example.com/auth to authorize, then retry." + fake := &fakeAuthenticator{ + outcome: &oauth.Outcome{UserAction: &oauth.UserAction{URL: "https://example.com/auth", Message: message}, FlowID: "flow-1"}, + } + + // No elicitation capability advertised, and no handler needed since the + // middleware should not issue an input request. + text, toolRan := runOAuthMiddlewareCall(t, fake, &mcp.ClientCapabilities{}, nil) + + assert.False(t, toolRan, "the tool must not run before authorization completes") + assert.Equal(t, message, text, "the manual instructions should be surfaced as a tool result") + assert.Zero(t, fake.awaitCalls) + assert.Zero(t, fake.cancelCalls) + }) +} + +func TestOAuthMultiRoundTripResultType(t *testing.T) { + t.Parallel() + + fake := &fakeAuthenticator{ + outcome: &oauth.Outcome{ + UserAction: &oauth.UserAction{URL: "https://example.com/auth", Message: "manual"}, + FlowID: "flow-1", + }, + } + server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "v0.0.1"}, nil) + var toolRan bool + handler := createOAuthToolMiddleware(fake, discardLogger())(func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + toolRan = true + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "tool-ran"}}}, nil + }) + server.AddTool(&mcp.Tool{ + Name: probeToolName, + InputSchema: &jsonschema.Schema{Type: "object"}, + }, handler) + + st, ct := mcp.NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = ss.Close() }) + + client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "v0.0.1"}, &mcp.ClientOptions{ + Capabilities: &mcp.ClientCapabilities{Elicitation: &mcp.ElicitationCapabilities{URL: &mcp.URLElicitationCapabilities{}}}, + MultiRoundTrip: &mcp.MultiRoundTripOptions{Disabled: true}, + }) + cs, err := client.Connect(context.Background(), ct, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = cs.Close() }) + + res, err := cs.CallTool(context.Background(), &mcp.CallToolParams{Name: probeToolName}) + require.NoError(t, err) + assert.True(t, res.NeedsInput(), "the wire response must declare resultType input_required") + assert.Contains(t, res.InputRequests, oauthElicitIDPrefix+"flow-1") + assert.False(t, toolRan) +} + // TestRunStdioServerRejectsTokenAndOAuth verifies the mutually-exclusive guard: // supplying both a static token and an OAuth manager is rejected before the // server starts, rather than silently preferring one for auth and the other for diff --git a/internal/ghmcp/server.go b/internal/ghmcp/server.go index 2267dd5d62..1e0611d648 100644 --- a/internal/ghmcp/server.go +++ b/internal/ghmcp/server.go @@ -314,38 +314,35 @@ func RunStdioServer(cfg StdioServerConfig) error { // For OAuth, the token is resolved lazily: empty until the user authorizes // on the first tool call, then refreshed for the rest of the session. var tokenProvider func() string + var toolHandlerMiddleware []inventory.ToolHandlerMiddleware if cfg.OAuthManager != nil { tokenProvider = cfg.OAuthManager.AccessToken + toolHandlerMiddleware = append(toolHandlerMiddleware, createOAuthToolMiddleware(cfg.OAuthManager, logger)) } ghServer, err := NewStdioMCPServer(ctx, github.MCPServerConfig{ - Version: cfg.Version, - Host: cfg.Host, - Token: cfg.Token, - EnabledToolsets: cfg.EnabledToolsets, - EnabledTools: cfg.EnabledTools, - EnabledFeatures: cfg.EnabledFeatures, - ReadOnly: cfg.ReadOnly, - Translator: t, - ContentWindowSize: cfg.ContentWindowSize, - LockdownMode: cfg.LockdownMode, - InsidersMode: cfg.InsidersMode, - ExcludeTools: cfg.ExcludeTools, - Logger: logger, - RepoAccessTTL: cfg.RepoAccessCacheTTL, - TokenScopes: tokenScopes, - TokenProvider: tokenProvider, + Version: cfg.Version, + Host: cfg.Host, + Token: cfg.Token, + EnabledToolsets: cfg.EnabledToolsets, + EnabledTools: cfg.EnabledTools, + EnabledFeatures: cfg.EnabledFeatures, + ReadOnly: cfg.ReadOnly, + Translator: t, + ContentWindowSize: cfg.ContentWindowSize, + LockdownMode: cfg.LockdownMode, + InsidersMode: cfg.InsidersMode, + ExcludeTools: cfg.ExcludeTools, + Logger: logger, + RepoAccessTTL: cfg.RepoAccessCacheTTL, + TokenScopes: tokenScopes, + TokenProvider: tokenProvider, + ToolHandlerMiddleware: toolHandlerMiddleware, }) if err != nil { return fmt.Errorf("failed to create MCP server: %w", err) } - // With OAuth, intercept tool calls to run the authorization flow on first - // use, before the handler tries to call GitHub with an empty token. - if cfg.OAuthManager != nil { - ghServer.AddReceivingMiddleware(createOAuthMiddleware(cfg.OAuthManager, logger)) - } - if cfg.ExportTranslations { // Once server is initialized, all translations are loaded dumpTranslations() diff --git a/internal/oauth/manager.go b/internal/oauth/manager.go index a78e919df7..8c16729f5b 100644 --- a/internal/oauth/manager.go +++ b/internal/oauth/manager.go @@ -2,6 +2,7 @@ package oauth import ( "context" + "crypto/rand" "errors" "log/slog" "net/http" @@ -20,6 +21,10 @@ const DefaultAuthTimeout = 5 * time.Minute // stalled GitHub token endpoint cannot block a tool call indefinitely. const tokenRefreshTimeout = 30 * time.Second +// ErrStaleAuthorizationFlow indicates that a prompt response belongs to an +// authorization flow that is no longer current. +var ErrStaleAuthorizationFlow = errors.New("authorization prompt has expired") + // flowStatus tracks the manager's single-flight authorization state. type flowStatus int @@ -37,6 +42,11 @@ type Outcome struct { // flow continues in the background; the user should retry once they have // completed it. UserAction *UserAction + + // FlowID correlates a user action with the authorization flow that produced + // it. Callers must pass it back to AwaitToken or Cancel so a delayed response + // cannot affect a newer flow. + FlowID string } // UserAction is an instruction for the user to complete authorization out of @@ -65,9 +75,12 @@ type Manager struct { mu sync.Mutex source oauth2.TokenSource // refreshing source, set once authorized + tokenGeneration uint64 // increments whenever source is replaced status flowStatus + flowID string pending *UserAction done chan struct{} + cancelFlow context.CancelFunc // cancels the in-flight flow, if any lastErr error refreshErrLogged bool // true once a refresh failure has been logged, reset on re-auth } @@ -93,11 +106,21 @@ func NewManager(cfg Config, logger *slog.Logger) *Manager { // re-authorization is required). It is cheap to call repeatedly: the underlying // token source caches and only refreshes when the token has expired. func (m *Manager) AccessToken() string { + token, _ := m.accessToken() + return token +} + +// accessToken returns the token together with the generation of the source it +// checked. Authenticate uses the generation to detect a source installed while +// token validation was in progress, without repeating a potentially blocking +// refresh request. +func (m *Manager) accessToken() (string, uint64) { m.mu.Lock() src := m.source + generation := m.tokenGeneration m.mu.Unlock() if src == nil { - return "" + return "", generation } // Refresh (if needed) happens here, off the lock, because ReuseTokenSource may // make a blocking network call and holding m.mu would serialize every tool call. @@ -109,17 +132,17 @@ func (m *Manager) AccessToken() string { // prompt. The oauth2 error carries the token endpoint's response, not the // access or refresh token. m.mu.Lock() - if !m.refreshErrLogged { + if m.tokenGeneration == generation && !m.refreshErrLogged { m.refreshErrLogged = true m.logger.Warn("OAuth token refresh failed; re-authorization required", "error", err) } m.mu.Unlock() - return "" + return "", generation } if !tok.Valid() { - return "" + return "", generation } - return tok.AccessToken + return tok.AccessToken, generation } // HasToken reports whether a valid token is currently available. @@ -137,63 +160,142 @@ func (m *Manager) HasToken() bool { // Only one flow runs at a time. Concurrent callers either join a running secure // flow, receive the pending user action, or are told to retry shortly. func (m *Manager) Authenticate(ctx context.Context, prompter Prompter) (*Outcome, error) { - if m.AccessToken() != "" { - return nil, nil - } + var flowID string + var done chan struct{} + for { + token, checkedTokenGeneration := m.accessToken() + if token != "" { + return nil, nil + } - m.mu.Lock() - switch m.status { - case statusAwaitingUser: - ua := m.pending - m.mu.Unlock() - return &Outcome{UserAction: ua}, nil - case statusStarting: - m.mu.Unlock() - return &Outcome{UserAction: &UserAction{ - Message: "GitHub authorization is already in progress. Please retry your request in a few seconds.", - }}, nil - case statusInProgress: - done := m.done + m.mu.Lock() + switch m.status { + case statusAwaitingUser: + ua := m.pending + flowID := m.flowID + m.mu.Unlock() + return &Outcome{UserAction: ua, FlowID: flowID}, nil + case statusStarting: + flowID := m.flowID + m.mu.Unlock() + return &Outcome{UserAction: &UserAction{ + Message: "GitHub authorization is already in progress. Please retry your request in a few seconds.", + }, FlowID: flowID}, nil + case statusInProgress: + done := m.done + flowID := m.flowID + m.mu.Unlock() + return m.joinWait(ctx, done, flowID) + } + + // A flow may have installed a token while the source above was being + // checked. Retry if the source changed before claiming the idle state. + if m.tokenGeneration != checkedTokenGeneration { + m.mu.Unlock() + continue + } + + // Idle: this call owns the new flow. + m.status = statusStarting + m.flowID = rand.Text() + flowID = m.flowID + m.lastErr = nil + m.done = make(chan struct{}) + done = m.done m.mu.Unlock() - return m.joinWait(ctx, done) + break } - // Idle: this call owns the new flow. - m.status = statusStarting - m.lastErr = nil - m.done = make(chan struct{}) - done := m.done - m.mu.Unlock() - plan, err := m.begin(prompter) if err != nil { - m.complete(nil, err) + m.complete(flowID, nil, err) return nil, err } + bgCtx, cancel := context.WithTimeout(context.Background(), DefaultAuthTimeout) m.mu.Lock() + if m.flowID != flowID { + m.mu.Unlock() + cancel() + return nil, ErrStaleAuthorizationFlow + } if plan.userAction != nil { m.status = statusAwaitingUser m.pending = plan.userAction } else { m.status = statusInProgress } + m.cancelFlow = cancel m.mu.Unlock() - - bgCtx, cancel := context.WithTimeout(context.Background(), DefaultAuthTimeout) - go m.runFlow(bgCtx, cancel, plan) + go m.runFlow(bgCtx, cancel, flowID, plan) if plan.userAction != nil { - return &Outcome{UserAction: plan.userAction}, nil + return &Outcome{UserAction: plan.userAction, FlowID: flowID}, nil } - return m.joinWait(ctx, done) + return m.joinWait(ctx, done, flowID) +} + +// AwaitToken blocks until the in-flight authorization flow yields a token, the +// flow ends without one, or ctx is done. It is the resume half of the +// multi-round-trip flow: a transport that presented the authorization prompt +// itself (via elicitation returned from a tool call) calls this once the user +// has acted, to wait for the background token acquisition to finish. +// +// It returns (nil, nil) once a token is available (proceed), (&Outcome{UserAction}, +// nil) when the user must still act out of band, or (nil, err) on failure. +func (m *Manager) AwaitToken(ctx context.Context, flowID string) (*Outcome, error) { + m.mu.Lock() + if flowID == "" || flowID != m.flowID { + m.mu.Unlock() + return nil, ErrStaleAuthorizationFlow + } + done := m.done + m.mu.Unlock() + if m.AccessToken() != "" { + return nil, nil + } + if done == nil { + // No flow is in flight; report whatever terminal state it left behind. + return m.outcomeAfterFlow(flowID) + } + select { + case <-done: + return m.outcomeAfterFlow(flowID) + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// Cancel retires the matching authorization flow and aborts its background +// callback listener or device poll. It returns false if flowID is stale. +func (m *Manager) Cancel(flowID string) bool { + m.mu.Lock() + if flowID == "" || flowID != m.flowID { + m.mu.Unlock() + return false + } + cancel := m.cancelFlow + m.status = statusIdle + m.flowID = "" + m.pending = nil + m.cancelFlow = nil + m.lastErr = context.Canceled + if m.done != nil { + close(m.done) + m.done = nil + } + m.mu.Unlock() + if cancel != nil { + cancel() + } + return true } // runFlow executes a prepared flow in the background and records the result. The // optional display prompt runs concurrently: a decline (or other failure) aborts // the flow, while an undeliverable prompt degrades to the manual fallback without // tearing the flow down, so the user can still authorize out of band. -func (m *Manager) runFlow(ctx context.Context, cancel context.CancelFunc, plan *flowPlan) { +func (m *Manager) runFlow(ctx context.Context, cancel context.CancelFunc, flowID string, plan *flowPlan) { defer cancel() if plan.display != nil { @@ -212,7 +314,7 @@ func (m *Manager) runFlow(ctx context.Context, cancel context.CancelFunc, plan * // prompt. Surface the manual instructions instead of failing, and // keep the background flow alive so the user can still authorize. m.logger.Debug("authorization prompt undeliverable; falling back to manual instructions", "reason", err) - m.fallBackToUserAction(plan.fallback) + m.fallBackToUserAction(flowID, plan.fallback) default: // A user decline (ErrPromptDeclined) or any other prompt failure // ends the flow. @@ -223,17 +325,17 @@ func (m *Manager) runFlow(ctx context.Context, cancel context.CancelFunc, plan * } tok, err := plan.run(ctx) - m.complete(tok, err) + m.complete(flowID, tok, err) } // fallBackToUserAction promotes a running secure flow to the manual user-action // channel after its prompt could not be delivered. The background flow keeps // running, so the user can complete authorization out of band and retry. It is a // no-op if the flow has already resolved. -func (m *Manager) fallBackToUserAction(ua *UserAction) { +func (m *Manager) fallBackToUserAction(flowID string, ua *UserAction) { m.mu.Lock() defer m.mu.Unlock() - if m.status != statusInProgress { + if m.flowID != flowID || m.status != statusInProgress { return } m.status = statusAwaitingUser @@ -248,12 +350,16 @@ func (m *Manager) fallBackToUserAction(ua *UserAction) { // complete records the flow result, installing a refreshing token source on // success, and wakes any joined callers. -func (m *Manager) complete(tok *oauth2.Token, err error) { +func (m *Manager) complete(flowID string, tok *oauth2.Token, err error) { m.mu.Lock() defer m.mu.Unlock() + if m.flowID != flowID { + return + } m.status = statusIdle m.pending = nil + m.cancelFlow = nil if err != nil { m.lastErr = err m.logger.Debug("oauth flow failed", "error", err) @@ -265,6 +371,7 @@ func (m *Manager) complete(tok *oauth2.Token, err error) { // client so a stalled token endpoint can't block a tool call forever. refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{Timeout: tokenRefreshTimeout}) m.source = m.refreshConfig.TokenSource(refreshCtx, tok) + m.tokenGeneration++ m.refreshErrLogged = false m.logger.Info("github authorization complete") } @@ -277,28 +384,39 @@ func (m *Manager) complete(tok *oauth2.Token, err error) { // joinWait blocks until the running flow finishes or ctx is cancelled. If the // flow was promoted to the manual channel while waiting (its prompt could not be // delivered), it returns that user action rather than an error. -func (m *Manager) joinWait(ctx context.Context, done chan struct{}) (*Outcome, error) { +func (m *Manager) joinWait(ctx context.Context, done chan struct{}, flowID string) (*Outcome, error) { select { case <-done: - if m.AccessToken() != "" { - return nil, nil - } - m.mu.Lock() - pending := m.pending - err := m.lastErr - m.mu.Unlock() - if pending != nil { - return &Outcome{UserAction: pending}, nil - } - if err != nil { - return nil, err - } - return nil, errors.New("authorization did not complete") + return m.outcomeAfterFlow(flowID) case <-ctx.Done(): return nil, ctx.Err() } } +// outcomeAfterFlow reports the result once the flow's done channel has closed +// (or when there is no flow in flight): a token to proceed (nil, nil), a pending +// user action to surface, or the flow's error. +func (m *Manager) outcomeAfterFlow(flowID string) (*Outcome, error) { + m.mu.Lock() + if flowID == "" || flowID != m.flowID { + m.mu.Unlock() + return nil, ErrStaleAuthorizationFlow + } + pending := m.pending + err := m.lastErr + m.mu.Unlock() + if m.AccessToken() != "" { + return nil, nil + } + if pending != nil { + return &Outcome{UserAction: pending, FlowID: flowID}, nil + } + if err != nil { + return nil, err + } + return nil, errors.New("authorization did not complete") +} + func (m *Manager) oauth2Config(redirectURL string) *oauth2.Config { return &oauth2.Config{ ClientID: m.config.ClientID, diff --git a/internal/oauth/manager_test.go b/internal/oauth/manager_test.go index 6f43c03ef9..52d5a54309 100644 --- a/internal/oauth/manager_test.go +++ b/internal/oauth/manager_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/oauth2" ) // newManager wires a Manager to the fake GitHub server. By default the browser @@ -167,6 +168,7 @@ func TestAuthenticateLastDitchUserAction(t *testing.T) { require.NoError(t, err) require.NotNil(t, out) require.NotNil(t, out.UserAction) + require.NotEmpty(t, out.FlowID) assert.NotEmpty(t, out.UserAction.URL) assert.Contains(t, out.UserAction.Message, "open this URL") assert.Contains(t, out.UserAction.Message, securityAdvisory, @@ -178,12 +180,127 @@ func TestAuthenticateLastDitchUserAction(t *testing.T) { require.NoError(t, err) require.NotNil(t, out2.UserAction) assert.Equal(t, out.UserAction.URL, out2.UserAction.URL) + assert.Equal(t, out.FlowID, out2.FlowID) // The user opens the URL out of band; the background flow then completes. require.NoError(t, browserGet(out.UserAction.URL)) assert.Equal(t, "gho_access", waitForToken(t, m)) } +func TestAwaitTokenCompletesCurrentFlow(t *testing.T) { + f := newFakeGitHub(t) + m := newManager(t, f) + m.openURL = func(string) error { return errors.New("no browser") } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + out, err := m.Authenticate(ctx, nil) + require.NoError(t, err) + require.NotNil(t, out) + require.NotNil(t, out.UserAction) + require.NotEmpty(t, out.FlowID) + + authDone := make(chan error, 1) + go func() { + authDone <- browserGet(out.UserAction.URL) + }() + + awaited, err := m.AwaitToken(ctx, out.FlowID) + require.NoError(t, err) + assert.Nil(t, awaited) + require.NoError(t, <-authDone) + assert.Equal(t, "gho_access", m.AccessToken()) +} + +func TestCancelAndAwaitTokenAreFlowScoped(t *testing.T) { + f := newFakeGitHub(t) + m := newManager(t, f) + m.openURL = func(string) error { return errors.New("no browser") } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + first, err := m.Authenticate(ctx, nil) + require.NoError(t, err) + require.NotNil(t, first) + require.NotEmpty(t, first.FlowID) + + assert.True(t, m.Cancel(first.FlowID), "the current flow should be cancelled") + second, err := m.Authenticate(ctx, nil) + require.NoError(t, err) + require.NotNil(t, second) + require.NotNil(t, second.UserAction) + require.NotEmpty(t, second.FlowID) + require.NotEqual(t, first.FlowID, second.FlowID) + + assert.False(t, m.Cancel(first.FlowID), "a stale decline must not cancel the newer flow") + _, err = m.AwaitToken(ctx, first.FlowID) + assert.ErrorIs(t, err, ErrStaleAuthorizationFlow) + + authDone := make(chan error, 1) + go func() { + authDone <- browserGet(second.UserAction.URL) + }() + awaited, err := m.AwaitToken(ctx, second.FlowID) + require.NoError(t, err) + assert.Nil(t, awaited) + require.NoError(t, <-authDone) + assert.Equal(t, "gho_access", m.AccessToken()) +} + +type blockingTokenSource struct { + entered chan struct{} + release chan struct{} + once sync.Once +} + +func (s *blockingTokenSource) Token() (*oauth2.Token, error) { + s.once.Do(func() { close(s.entered) }) + <-s.release + return nil, errors.New("stale token source failed") +} + +func TestAuthenticateRechecksTokenBeforeStartingFlow(t *testing.T) { + f := newFakeGitHub(t) + m := newManager(t, f) + staleSource := &blockingTokenSource{ + entered: make(chan struct{}), + release: make(chan struct{}), + } + + m.mu.Lock() + m.source = staleSource + m.status = statusInProgress + m.flowID = "existing-flow" + m.done = make(chan struct{}) + m.mu.Unlock() + + type result struct { + out *Outcome + err error + } + authResult := make(chan result, 1) + go func() { + out, err := m.Authenticate(context.Background(), nil) + authResult <- result{out: out, err: err} + }() + + <-staleSource.entered + m.complete("existing-flow", &oauth2.Token{ + AccessToken: "installed-token", + TokenType: "bearer", + Expiry: time.Now().Add(time.Hour), + }, nil) + close(staleSource.release) + + got := <-authResult + require.NoError(t, got.err) + assert.Nil(t, got.out) + assert.Equal(t, "installed-token", m.AccessToken()) + assert.Empty(t, f.recordedGrants(), "a completed concurrent flow must not trigger a redundant authorization") +} + func TestAuthenticateDeviceFlow(t *testing.T) { f := newFakeGitHub(t) f.deviceToken = "gho_device_token" diff --git a/pkg/github/server.go b/pkg/github/server.go index 627cc678b2..67db83a77a 100644 --- a/pkg/github/server.go +++ b/pkg/github/server.go @@ -73,6 +73,11 @@ type MCPServerConfig struct { // token is obtained lazily on first use and refreshed thereafter. TokenProvider func() string + // ToolHandlerMiddleware wraps every registered tool handler. Unlike MCP + // receiving middleware, these wrappers execute inside Server.callTool, so + // SDK result finalization still runs on results they return. + ToolHandlerMiddleware []inventory.ToolHandlerMiddleware + // Additional server options to apply ServerOptions []MCPServerOption } @@ -105,7 +110,7 @@ func NewMCPServer(ctx context.Context, cfg *MCPServerConfig, deps ToolDependenci } // Register GitHub tools/resources/prompts from the inventory. - inv.RegisterAll(ctx, ghServer, deps) + inv.RegisterAll(ctx, ghServer, deps, cfg.ToolHandlerMiddleware...) // Register MCP App UI resources whenever the embedded UI assets are // available. The resources are static HTML and are only referenced by diff --git a/pkg/inventory/registry.go b/pkg/inventory/registry.go index 6505e6b5ef..915ed0aa1c 100644 --- a/pkg/inventory/registry.go +++ b/pkg/inventory/registry.go @@ -219,9 +219,9 @@ func shouldStripMCPAppsMetadata(ctx context.Context, featureFlagEnabled bool) bo // user identity from ctx would otherwise see context.Background() and // falsely report the flag off, even when the actual request arrived on the // /insiders route. -func (r *Inventory) RegisterTools(ctx context.Context, s *mcp.Server, deps any) { +func (r *Inventory) RegisterTools(ctx context.Context, s *mcp.Server, deps any, middleware ...ToolHandlerMiddleware) { for _, tool := range r.ToolsForRegistration(ctx) { - tool.RegisterFunc(s, deps) + tool.RegisterFunc(s, deps, middleware...) } } @@ -257,8 +257,8 @@ func (r *Inventory) RegisterPrompts(ctx context.Context, s *mcp.Server) { // RegisterAll registers all available tools, resources, and prompts with the server. // The context is used for feature flag evaluation. -func (r *Inventory) RegisterAll(ctx context.Context, s *mcp.Server, deps any) { - r.RegisterTools(ctx, s, deps) +func (r *Inventory) RegisterAll(ctx context.Context, s *mcp.Server, deps any, middleware ...ToolHandlerMiddleware) { + r.RegisterTools(ctx, s, deps, middleware...) r.RegisterResourceTemplates(ctx, s, deps) r.RegisterPrompts(ctx, s) } diff --git a/pkg/inventory/server_tool.go b/pkg/inventory/server_tool.go index beb70138eb..44a062ba2e 100644 --- a/pkg/inventory/server_tool.go +++ b/pkg/inventory/server_tool.go @@ -18,6 +18,10 @@ import ( // should define their own typed dependencies struct and type-assert as needed. type HandlerFunc func(deps any) mcp.ToolHandler +// ToolHandlerMiddleware wraps an MCP tool handler. Middleware is applied from +// right to left, so the first middleware passed to RegisterFunc executes first. +type ToolHandlerMiddleware func(next mcp.ToolHandler) mcp.ToolHandler + // ToolsetID is a unique identifier for a toolset. // Using a distinct type provides compile-time type safety. type ToolsetID string @@ -110,8 +114,11 @@ func (st *ServerTool) Handler(deps any) mcp.ToolHandler { // Icons are automatically applied from the toolset metadata if not already set. // A shallow copy of the tool is made to avoid mutating the original ServerTool. // Panics if the tool has no handler - all tools should have handlers. -func (st *ServerTool) RegisterFunc(s *mcp.Server, deps any) { +func (st *ServerTool) RegisterFunc(s *mcp.Server, deps any, middleware ...ToolHandlerMiddleware) { handler := st.Handler(deps) // This will panic if HandlerFunc is nil + for i := len(middleware) - 1; i >= 0; i-- { + handler = middleware[i](handler) + } // Make a shallow copy of the tool to avoid mutating the original toolCopy := st.Tool // Apply icons from toolset metadata if tool doesn't have icons set diff --git a/pkg/inventory/server_tool_test.go b/pkg/inventory/server_tool_test.go index adf012b1f5..c6d2a6fdd8 100644 --- a/pkg/inventory/server_tool_test.go +++ b/pkg/inventory/server_tool_test.go @@ -81,6 +81,51 @@ func TestNewServerToolWithContextHandler_ValidArguments_Succeeds(t *testing.T) { assert.Equal(t, "success: octocat/hello-world", textContent.Text) } +func TestServerToolRegisterFuncAppliesMiddleware(t *testing.T) { + tool := NewServerTool( + mcp.Tool{ + Name: "wrapped_tool", + InputSchema: &jsonschema.Schema{Type: "object"}, + }, + testToolsetMetadata("test"), + func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "handler"}}, + }, nil + }, + ) + + middlewareCalled := make(chan struct{}, 1) + middleware := func(next mcp.ToolHandler) mcp.ToolHandler { + return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + middlewareCalled <- struct{}{} + return next(ctx, req) + } + } + + server := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "v0.0.1"}, nil) + tool.RegisterFunc(server, nil, middleware) + st, ct := mcp.NewInMemoryTransports() + ss, err := server.Connect(context.Background(), st, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = ss.Close() }) + + client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "v0.0.1"}, nil) + cs, err := client.Connect(context.Background(), ct, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = cs.Close() }) + + result, err := cs.CallTool(context.Background(), &mcp.CallToolParams{Name: "wrapped_tool"}) + require.NoError(t, err) + select { + case <-middlewareCalled: + default: + t.Fatal("tool middleware was not called") + } + require.Len(t, result.Content, 1) + assert.Equal(t, "handler", result.Content[0].(*mcp.TextContent).Text) +} + func TestAnnotateHeaderParams(t *testing.T) { tool := &mcp.Tool{InputSchema: &jsonschema.Schema{ Type: "object", diff --git a/third-party-licenses.darwin.md b/third-party-licenses.darwin.md index 88235f3f40..fca63f2256 100644 --- a/third-party-licenses.darwin.md +++ b/third-party-licenses.darwin.md @@ -24,8 +24,8 @@ The following packages are included for the amd64, arm64 architectures. - [github.com/josephburnett/jd/v2](https://pkg.go.dev/github.com/josephburnett/jd/v2) ([MIT](https://github.com/josephburnett/jd/blob/v2.5.0/v2/LICENSE)) - [github.com/lithammer/fuzzysearch/fuzzy](https://pkg.go.dev/github.com/lithammer/fuzzysearch/fuzzy) ([MIT](https://github.com/lithammer/fuzzysearch/blob/v1.1.8/LICENSE)) - [github.com/microcosm-cc/bluemonday](https://pkg.go.dev/github.com/microcosm-cc/bluemonday) ([BSD-3-Clause](https://github.com/microcosm-cc/bluemonday/blob/v1.0.27/LICENSE.md)) - - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([Apache-2.0](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.1/LICENSE)) - - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([MIT](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.1/LICENSE)) + - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([Apache-2.0](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.2/LICENSE)) + - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([MIT](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.2/LICENSE)) - [github.com/muesli/cache2go](https://pkg.go.dev/github.com/muesli/cache2go) ([BSD-3-Clause](https://github.com/muesli/cache2go/blob/518229cd8021/LICENSE.txt)) - [github.com/pelletier/go-toml/v2](https://pkg.go.dev/github.com/pelletier/go-toml/v2) ([MIT](https://github.com/pelletier/go-toml/blob/v2.2.4/LICENSE)) - [github.com/sagikazarmark/locafero](https://pkg.go.dev/github.com/sagikazarmark/locafero) ([MIT](https://github.com/sagikazarmark/locafero/blob/v0.11.0/LICENSE)) diff --git a/third-party-licenses.linux.md b/third-party-licenses.linux.md index e3762f5c04..dc3798c769 100644 --- a/third-party-licenses.linux.md +++ b/third-party-licenses.linux.md @@ -24,8 +24,8 @@ The following packages are included for the 386, amd64, arm64 architectures. - [github.com/josephburnett/jd/v2](https://pkg.go.dev/github.com/josephburnett/jd/v2) ([MIT](https://github.com/josephburnett/jd/blob/v2.5.0/v2/LICENSE)) - [github.com/lithammer/fuzzysearch/fuzzy](https://pkg.go.dev/github.com/lithammer/fuzzysearch/fuzzy) ([MIT](https://github.com/lithammer/fuzzysearch/blob/v1.1.8/LICENSE)) - [github.com/microcosm-cc/bluemonday](https://pkg.go.dev/github.com/microcosm-cc/bluemonday) ([BSD-3-Clause](https://github.com/microcosm-cc/bluemonday/blob/v1.0.27/LICENSE.md)) - - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([Apache-2.0](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.1/LICENSE)) - - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([MIT](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.1/LICENSE)) + - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([Apache-2.0](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.2/LICENSE)) + - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([MIT](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.2/LICENSE)) - [github.com/muesli/cache2go](https://pkg.go.dev/github.com/muesli/cache2go) ([BSD-3-Clause](https://github.com/muesli/cache2go/blob/518229cd8021/LICENSE.txt)) - [github.com/pelletier/go-toml/v2](https://pkg.go.dev/github.com/pelletier/go-toml/v2) ([MIT](https://github.com/pelletier/go-toml/blob/v2.2.4/LICENSE)) - [github.com/sagikazarmark/locafero](https://pkg.go.dev/github.com/sagikazarmark/locafero) ([MIT](https://github.com/sagikazarmark/locafero/blob/v0.11.0/LICENSE)) diff --git a/third-party-licenses.windows.md b/third-party-licenses.windows.md index eb0743558a..db7db1ec1b 100644 --- a/third-party-licenses.windows.md +++ b/third-party-licenses.windows.md @@ -25,8 +25,8 @@ The following packages are included for the 386, amd64, arm64 architectures. - [github.com/josephburnett/jd/v2](https://pkg.go.dev/github.com/josephburnett/jd/v2) ([MIT](https://github.com/josephburnett/jd/blob/v2.5.0/v2/LICENSE)) - [github.com/lithammer/fuzzysearch/fuzzy](https://pkg.go.dev/github.com/lithammer/fuzzysearch/fuzzy) ([MIT](https://github.com/lithammer/fuzzysearch/blob/v1.1.8/LICENSE)) - [github.com/microcosm-cc/bluemonday](https://pkg.go.dev/github.com/microcosm-cc/bluemonday) ([BSD-3-Clause](https://github.com/microcosm-cc/bluemonday/blob/v1.0.27/LICENSE.md)) - - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([Apache-2.0](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.1/LICENSE)) - - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([MIT](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.1/LICENSE)) + - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([Apache-2.0](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.2/LICENSE)) + - [github.com/modelcontextprotocol/go-sdk](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) ([MIT](https://github.com/modelcontextprotocol/go-sdk/blob/v1.7.0-pre.2/LICENSE)) - [github.com/muesli/cache2go](https://pkg.go.dev/github.com/muesli/cache2go) ([BSD-3-Clause](https://github.com/muesli/cache2go/blob/518229cd8021/LICENSE.txt)) - [github.com/pelletier/go-toml/v2](https://pkg.go.dev/github.com/pelletier/go-toml/v2) ([MIT](https://github.com/pelletier/go-toml/blob/v2.2.4/LICENSE)) - [github.com/sagikazarmark/locafero](https://pkg.go.dev/github.com/sagikazarmark/locafero) ([MIT](https://github.com/sagikazarmark/locafero/blob/v0.11.0/LICENSE))