mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(coderd): refresh expired MCP OAuth2 tokens everywhere (#23713)
Fixes expired MCP OAuth2 tokens causing 401 errors and stale `auth_connected` status in the UI. When users authenticate MCP servers (e.g. GitHub) via OAuth2, the access token and refresh token are stored in the database. However, when the access token expired, nothing refreshed it anywhere: - **chatd**: sent the expired token as-is, getting a 401 and skipping the MCP server - **list/get endpoints**: reported `auth_connected: true` just because a token record existed, regardless of expiry ## Changes ### Shared utility: `mcpclient.RefreshOAuth2Token` Pure function that uses `golang.org/x/oauth2` `TokenSource` to check if a token is expired (or within 10s of expiry) and refresh it. No DB dependency — callers handle persistence. ### chatd (`coderd/x/chatd/chatd.go`) Before calling `mcpclient.ConnectAll`, refreshes expired tokens. Persists new credentials to the database. Falls back to the old token if refresh fails. ### List/get MCP server endpoints (`coderd/mcp.go`) Both `listMCPServerConfigs` and `getMCPServerConfig` now attempt refresh when checking `auth_connected`. If the token is expired: - **Has refresh token**: attempt refresh, persist result, report `auth_connected` based on success - **No refresh token**: report `auth_connected: false` if expired This means the UI accurately reflects whether the user's token is actually usable, rather than just whether a record exists. <details> <summary>Design notes</summary> - `RefreshOAuth2Token` lives in `mcpclient` to avoid circular imports (`coderd` → `chatd` → `mcpclient` is fine; `chatd` → `coderd` would be circular). - DB persistence is handled by each caller with their own authz context (`AsSystemRestricted` in both cases). - The `buildAuthHeaders` warning in mcpclient about expired tokens is kept as defense-in-depth logging. </details>
This commit is contained in:
+82
-7
@@ -25,6 +25,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/httpmw"
|
||||
"github.com/coder/coder/v2/coderd/rbac"
|
||||
"github.com/coder/coder/v2/coderd/rbac/policy"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/mcpclient"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
@@ -59,7 +60,8 @@ func (api *API) listMCPServerConfigs(rw http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
// Look up the calling user's OAuth2 tokens so we can populate
|
||||
// auth_connected per server.
|
||||
// auth_connected per server. Attempt to refresh expired tokens
|
||||
// so the status is accurate and the token is ready for use.
|
||||
//nolint:gocritic // Need to check user tokens across all servers.
|
||||
userTokens, err := api.Database.GetMCPServerUserTokensByUserID(dbauthz.AsSystemRestricted(ctx), apiKey.UserID)
|
||||
if err != nil {
|
||||
@@ -69,9 +71,20 @@ func (api *API) listMCPServerConfigs(rw http.ResponseWriter, r *http.Request) {
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Build a config lookup for the refresh helper.
|
||||
configByID := make(map[uuid.UUID]database.MCPServerConfig, len(configs))
|
||||
for _, c := range configs {
|
||||
configByID[c.ID] = c
|
||||
}
|
||||
|
||||
tokenMap := make(map[uuid.UUID]bool, len(userTokens))
|
||||
for _, t := range userTokens {
|
||||
tokenMap[t.MCPServerConfigID] = true
|
||||
for _, tok := range userTokens {
|
||||
cfg, ok := configByID[tok.MCPServerConfigID]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
tokenMap[tok.MCPServerConfigID] = api.refreshMCPUserToken(ctx, cfg, tok, apiKey.UserID)
|
||||
}
|
||||
|
||||
resp := make([]codersdk.MCPServerConfig, 0, len(configs))
|
||||
@@ -386,7 +399,8 @@ func (api *API) getMCPServerConfig(rw http.ResponseWriter, r *http.Request) {
|
||||
sdkConfig = convertMCPServerConfigRedacted(config)
|
||||
}
|
||||
|
||||
// Populate AuthConnected for the calling user.
|
||||
// Populate AuthConnected for the calling user. Attempt to
|
||||
// refresh the token so the status is accurate.
|
||||
if config.AuthType == "oauth2" {
|
||||
//nolint:gocritic // Need to check user token for this server.
|
||||
userTokens, err := api.Database.GetMCPServerUserTokensByUserID(dbauthz.AsSystemRestricted(ctx), apiKey.UserID)
|
||||
@@ -397,9 +411,9 @@ func (api *API) getMCPServerConfig(rw http.ResponseWriter, r *http.Request) {
|
||||
})
|
||||
return
|
||||
}
|
||||
for _, t := range userTokens {
|
||||
if t.MCPServerConfigID == config.ID {
|
||||
sdkConfig.AuthConnected = true
|
||||
for _, tok := range userTokens {
|
||||
if tok.MCPServerConfigID == config.ID {
|
||||
sdkConfig.AuthConnected = api.refreshMCPUserToken(ctx, config, tok, apiKey.UserID)
|
||||
break
|
||||
}
|
||||
}
|
||||
@@ -1002,6 +1016,67 @@ func (api *API) mcpServerOAuth2Disconnect(rw http.ResponseWriter, r *http.Reques
|
||||
|
||||
// parseMCPServerConfigID extracts the MCP server config UUID from the
|
||||
// "mcpServer" path parameter.
|
||||
// refreshMCPUserToken attempts to refresh an expired OAuth2 token
|
||||
// for the given MCP server config. Returns true when the token is
|
||||
// valid (either still fresh or successfully refreshed), false when
|
||||
// the token is expired and cannot be refreshed.
|
||||
func (api *API) refreshMCPUserToken(
|
||||
ctx context.Context,
|
||||
cfg database.MCPServerConfig,
|
||||
tok database.MCPServerUserToken,
|
||||
userID uuid.UUID,
|
||||
) bool {
|
||||
if cfg.AuthType != "oauth2" {
|
||||
return true
|
||||
}
|
||||
if tok.RefreshToken == "" {
|
||||
// No refresh token — consider connected only if not
|
||||
// expired (or no expiry set).
|
||||
return !tok.Expiry.Valid || tok.Expiry.Time.After(time.Now())
|
||||
}
|
||||
|
||||
result, err := mcpclient.RefreshOAuth2Token(ctx, cfg, tok)
|
||||
if err != nil {
|
||||
api.Logger.Warn(ctx, "failed to refresh MCP oauth2 token",
|
||||
slog.F("server_slug", cfg.Slug),
|
||||
slog.Error(err),
|
||||
)
|
||||
// Refresh failed — token is dead.
|
||||
return false
|
||||
}
|
||||
|
||||
if result.Refreshed {
|
||||
var expiry sql.NullTime
|
||||
if !result.Expiry.IsZero() {
|
||||
expiry = sql.NullTime{Time: result.Expiry, Valid: true}
|
||||
}
|
||||
|
||||
//nolint:gocritic // Need system-level write access to
|
||||
// persist the refreshed OAuth2 token.
|
||||
_, err = api.Database.UpsertMCPServerUserToken(
|
||||
dbauthz.AsSystemRestricted(ctx),
|
||||
database.UpsertMCPServerUserTokenParams{
|
||||
MCPServerConfigID: tok.MCPServerConfigID,
|
||||
UserID: userID,
|
||||
AccessToken: result.AccessToken,
|
||||
AccessTokenKeyID: sql.NullString{},
|
||||
RefreshToken: result.RefreshToken,
|
||||
RefreshTokenKeyID: sql.NullString{},
|
||||
TokenType: result.TokenType,
|
||||
Expiry: expiry,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
api.Logger.Warn(ctx, "failed to persist refreshed MCP oauth2 token",
|
||||
slog.F("server_slug", cfg.Slug),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func parseMCPServerConfigID(rw http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
|
||||
mcpServerID, err := uuid.Parse(chi.URLParam(r, "mcpServer"))
|
||||
if err != nil {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -3857,6 +3858,8 @@ func (p *Server) runChat(
|
||||
})
|
||||
if len(mcpConfigs) > 0 {
|
||||
g2.Go(func() error {
|
||||
// Refresh expired OAuth2 tokens before connecting.
|
||||
mcpTokens = p.refreshExpiredMCPTokens(ctx, logger, mcpConfigs, mcpTokens)
|
||||
mcpTools, mcpCleanup = mcpclient.ConnectAll(
|
||||
ctx, logger, mcpConfigs, mcpTokens,
|
||||
)
|
||||
@@ -5107,3 +5110,111 @@ func (p *Server) Close() error {
|
||||
p.inflight.Wait()
|
||||
return nil
|
||||
}
|
||||
|
||||
// refreshExpiredMCPTokens checks each MCP OAuth2 token and refreshes
|
||||
// any that are expired (or about to expire). Tokens without a
|
||||
// refresh_token or that fail to refresh are returned unchanged so the
|
||||
// caller can still attempt the connection (which will likely fail with
|
||||
// a 401 for the expired ones).
|
||||
func (p *Server) refreshExpiredMCPTokens(
|
||||
ctx context.Context,
|
||||
logger slog.Logger,
|
||||
configs []database.MCPServerConfig,
|
||||
tokens []database.MCPServerUserToken,
|
||||
) []database.MCPServerUserToken {
|
||||
configsByID := make(map[uuid.UUID]database.MCPServerConfig, len(configs))
|
||||
for _, cfg := range configs {
|
||||
configsByID[cfg.ID] = cfg
|
||||
}
|
||||
|
||||
result := slices.Clone(tokens)
|
||||
|
||||
var eg errgroup.Group
|
||||
for i, tok := range result {
|
||||
cfg, ok := configsByID[tok.MCPServerConfigID]
|
||||
if !ok || cfg.AuthType != "oauth2" {
|
||||
continue
|
||||
}
|
||||
if tok.RefreshToken == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
eg.Go(func() error {
|
||||
refreshed, err := p.refreshMCPTokenIfNeeded(ctx, logger, cfg, tok)
|
||||
if err != nil {
|
||||
logger.Warn(ctx, "failed to refresh MCP oauth2 token",
|
||||
slog.F("server_slug", cfg.Slug),
|
||||
slog.Error(err),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
result[i] = refreshed
|
||||
return nil
|
||||
})
|
||||
}
|
||||
_ = eg.Wait()
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// refreshMCPTokenIfNeeded delegates to mcpclient.RefreshOAuth2Token
|
||||
// and persists the result to the database when a refresh occurs.
|
||||
// The logger should carry chat-scoped fields so log lines can be
|
||||
// correlated with specific chat requests.
|
||||
func (p *Server) refreshMCPTokenIfNeeded(
|
||||
ctx context.Context,
|
||||
logger slog.Logger,
|
||||
cfg database.MCPServerConfig,
|
||||
tok database.MCPServerUserToken,
|
||||
) (database.MCPServerUserToken, error) {
|
||||
result, err := mcpclient.RefreshOAuth2Token(ctx, cfg, tok)
|
||||
if err != nil {
|
||||
return tok, err
|
||||
}
|
||||
|
||||
if !result.Refreshed {
|
||||
return tok, nil
|
||||
}
|
||||
|
||||
logger.Info(ctx, "refreshed MCP oauth2 token",
|
||||
slog.F("server_slug", cfg.Slug),
|
||||
slog.F("user_id", tok.UserID),
|
||||
)
|
||||
|
||||
var expiry sql.NullTime
|
||||
if !result.Expiry.IsZero() {
|
||||
expiry = sql.NullTime{Time: result.Expiry, Valid: true}
|
||||
}
|
||||
|
||||
//nolint:gocritic // Chatd needs system-level write access to
|
||||
// persist the refreshed OAuth2 token for the user.
|
||||
updated, err := p.db.UpsertMCPServerUserToken(
|
||||
dbauthz.AsSystemRestricted(ctx),
|
||||
database.UpsertMCPServerUserTokenParams{
|
||||
MCPServerConfigID: tok.MCPServerConfigID,
|
||||
UserID: tok.UserID,
|
||||
AccessToken: result.AccessToken,
|
||||
AccessTokenKeyID: sql.NullString{},
|
||||
RefreshToken: result.RefreshToken,
|
||||
RefreshTokenKeyID: sql.NullString{},
|
||||
TokenType: result.TokenType,
|
||||
Expiry: expiry,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
// The provider may have rotated the refresh token,
|
||||
// invalidating the old one. Use the new token
|
||||
// in-memory so at least this connection succeeds.
|
||||
logger.Warn(ctx, "failed to persist refreshed MCP oauth2 token, using in-memory",
|
||||
slog.F("server_slug", cfg.Slug),
|
||||
slog.Error(err),
|
||||
)
|
||||
tok.AccessToken = result.AccessToken
|
||||
tok.RefreshToken = result.RefreshToken
|
||||
tok.TokenType = result.TokenType
|
||||
tok.Expiry = expiry
|
||||
return tok, nil
|
||||
}
|
||||
|
||||
return updated, nil
|
||||
}
|
||||
|
||||
@@ -3853,6 +3853,324 @@ func TestMCPServerToolInvocation(t *testing.T) {
|
||||
"MCP tool result should be persisted as a tool message in the database")
|
||||
}
|
||||
|
||||
// TestMCPServerOAuth2TokenRefresh verifies that when a chat uses an
|
||||
// MCP server with OAuth2 auth and the stored access token is expired,
|
||||
// chatd refreshes the token using the stored refresh_token before
|
||||
// connecting. The refreshed token is persisted to the database and
|
||||
// the MCP tool call succeeds.
|
||||
func TestMCPServerOAuth2TokenRefresh(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
// The "fresh" token that the mock OAuth2 server returns after
|
||||
// a successful refresh_token grant.
|
||||
freshAccessToken := "fresh-access-token-" + uuid.New().String()
|
||||
|
||||
// Mock OAuth2 token endpoint that exchanges a refresh token
|
||||
// for a new access token.
|
||||
var refreshCalled atomic.Int32
|
||||
tokenSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
refreshCalled.Add(1)
|
||||
|
||||
if r.Method != http.MethodPost {
|
||||
w.WriteHeader(http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
|
||||
grantType := r.FormValue("grant_type")
|
||||
if grantType != "refresh_token" {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_, _ = w.Write([]byte(`{"error":"unsupported_grant_type"}`))
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = fmt.Fprintf(w, `{"access_token":%q,"token_type":"Bearer","expires_in":3600,"refresh_token":"rotated-refresh-token"}`, freshAccessToken)
|
||||
}))
|
||||
t.Cleanup(tokenSrv.Close)
|
||||
|
||||
// Start a real MCP server with an auth middleware that only
|
||||
// accepts the fresh access token. An expired token (or any
|
||||
// other value) gets a 401.
|
||||
mcpSrv := mcpserver.NewMCPServer("authed-mcp", "1.0.0")
|
||||
mcpSrv.AddTools(mcpserver.ServerTool{
|
||||
Tool: mcpgo.NewTool("echo",
|
||||
mcpgo.WithDescription("Echoes the input"),
|
||||
mcpgo.WithString("input",
|
||||
mcpgo.Description("The input string"),
|
||||
mcpgo.Required(),
|
||||
),
|
||||
),
|
||||
Handler: func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
|
||||
input, _ := req.GetArguments()["input"].(string)
|
||||
return mcpgo.NewToolResultText("echo: " + input), nil
|
||||
},
|
||||
})
|
||||
mcpHTTP := mcpserver.NewStreamableHTTPServer(mcpSrv)
|
||||
// Wrap with auth check.
|
||||
authMux := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
auth := r.Header.Get("Authorization")
|
||||
if auth != "Bearer "+freshAccessToken {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"error":"invalid_token","error_description":"The access token is invalid or expired"}`))
|
||||
return
|
||||
}
|
||||
mcpHTTP.ServeHTTP(w, r)
|
||||
})
|
||||
mcpTS := httptest.NewServer(authMux)
|
||||
t.Cleanup(mcpTS.Close)
|
||||
|
||||
// Track LLM interactions.
|
||||
var (
|
||||
callCount atomic.Int32
|
||||
llmToolNames []string
|
||||
llmToolsMu sync.Mutex
|
||||
foundMCPResult atomic.Bool
|
||||
)
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
if callCount.Add(1) == 1 {
|
||||
names := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
names = append(names, tool.Function.Name)
|
||||
}
|
||||
llmToolsMu.Lock()
|
||||
llmToolNames = names
|
||||
llmToolsMu.Unlock()
|
||||
|
||||
// Ask the LLM to call the MCP echo tool.
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAIToolCallChunk(
|
||||
"authed-mcp__echo",
|
||||
`{"input":"hello via refreshed token"}`,
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
// Second call: verify the tool result was fed back.
|
||||
for _, msg := range req.Messages {
|
||||
if msg.Role == "tool" && strings.Contains(msg.Content, "echo: hello via refreshed token") {
|
||||
foundMCPResult.Store(true)
|
||||
}
|
||||
}
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("Done!")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
|
||||
// Seed the MCP server config with OAuth2 auth pointing to our
|
||||
// mock token endpoint.
|
||||
mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||
DisplayName: "Authed MCP",
|
||||
Slug: "authed-mcp",
|
||||
Url: mcpTS.URL,
|
||||
Transport: "streamable_http",
|
||||
AuthType: "oauth2",
|
||||
OAuth2ClientID: "test-client-id",
|
||||
OAuth2TokenURL: tokenSrv.URL,
|
||||
Availability: "default_off",
|
||||
Enabled: true,
|
||||
ToolAllowList: []string{},
|
||||
ToolDenyList: []string{},
|
||||
CreatedBy: user.ID,
|
||||
UpdatedBy: user.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Seed an expired OAuth2 token with a valid refresh_token.
|
||||
_, err = db.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{
|
||||
MCPServerConfigID: mcpConfig.ID,
|
||||
UserID: user.ID,
|
||||
AccessToken: "old-expired-access-token",
|
||||
RefreshToken: "old-refresh-token",
|
||||
TokenType: "Bearer",
|
||||
Expiry: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes()
|
||||
mockConn.EXPECT().ListMCPTools(gomock.Any()).
|
||||
Return(workspacesdk.ListMCPToolsResponse{}, nil).AnyTimes()
|
||||
mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.LSResponse{}, nil).AnyTimes()
|
||||
mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes()
|
||||
|
||||
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
|
||||
cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
require.Equal(t, dbAgent.ID, agentID)
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
})
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "oauth2-refresh-test",
|
||||
ModelConfigID: model.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
|
||||
MCPServerIDs: []uuid.UUID{mcpConfig.ID},
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("Echo something via the authed MCP."),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for the chat to finish processing.
|
||||
var chatResult database.Chat
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
chatResult = got
|
||||
return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
if chatResult.Status == database.ChatStatusError {
|
||||
require.FailNowf(t, "chat failed", "last_error=%q", chatResult.LastError.String)
|
||||
}
|
||||
|
||||
// The token should have been refreshed.
|
||||
require.Greater(t, refreshCalled.Load(), int32(0),
|
||||
"OAuth2 token endpoint should have been called to refresh the expired token")
|
||||
|
||||
// The MCP tool should appear in the tool list.
|
||||
llmToolsMu.Lock()
|
||||
recordedNames := append([]string(nil), llmToolNames...)
|
||||
llmToolsMu.Unlock()
|
||||
require.Contains(t, recordedNames, "authed-mcp__echo",
|
||||
"MCP tool should be in the tool list sent to the LLM")
|
||||
|
||||
// The tool result should have been fed back to the LLM.
|
||||
require.True(t, foundMCPResult.Load(),
|
||||
"MCP tool result should appear in the second LLM call")
|
||||
|
||||
// Verify the refreshed token was persisted to the database.
|
||||
dbToken, err := db.GetMCPServerUserToken(ctx, database.GetMCPServerUserTokenParams{
|
||||
MCPServerConfigID: mcpConfig.ID,
|
||||
UserID: user.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, freshAccessToken, dbToken.AccessToken,
|
||||
"refreshed access token should be persisted in the database")
|
||||
require.Equal(t, "rotated-refresh-token", dbToken.RefreshToken,
|
||||
"rotated refresh token should be persisted in the database")
|
||||
}
|
||||
|
||||
// TestMCPServerOAuth2TokenRefreshFailureGraceful verifies that when
|
||||
// the OAuth2 token endpoint is down, the chat still proceeds without
|
||||
// the MCP server's tools. The expired token is preserved unchanged.
|
||||
func TestMCPServerOAuth2TokenRefreshFailureGraceful(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
// Token endpoint that always returns an error.
|
||||
tokenSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
_, _ = w.Write([]byte(`{"error":"server_error","error_description":"token endpoint unavailable"}`))
|
||||
}))
|
||||
t.Cleanup(tokenSrv.Close)
|
||||
|
||||
// The LLM just replies with text — no tool calls.
|
||||
var callCount atomic.Int32
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
callCount.Add(1)
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("I responded without MCP tools.")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
|
||||
mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||
DisplayName: "Broken MCP",
|
||||
Slug: "broken-mcp",
|
||||
Url: "http://127.0.0.1:0/does-not-exist",
|
||||
Transport: "streamable_http",
|
||||
AuthType: "oauth2",
|
||||
OAuth2ClientID: "test-client-id",
|
||||
OAuth2TokenURL: tokenSrv.URL,
|
||||
Availability: "default_off",
|
||||
Enabled: true,
|
||||
ToolAllowList: []string{},
|
||||
ToolDenyList: []string{},
|
||||
CreatedBy: user.ID,
|
||||
UpdatedBy: user.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = db.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{
|
||||
MCPServerConfigID: mcpConfig.ID,
|
||||
UserID: user.ID,
|
||||
AccessToken: "old-expired-token",
|
||||
RefreshToken: "old-refresh-token",
|
||||
TokenType: "Bearer",
|
||||
Expiry: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "graceful-degradation-test",
|
||||
ModelConfigID: model.ID,
|
||||
MCPServerIDs: []uuid.UUID{mcpConfig.ID},
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("Hello, just reply."),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Chat should finish successfully despite the failed refresh.
|
||||
var chatResult database.Chat
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
chatResult = got
|
||||
return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
if chatResult.Status == database.ChatStatusError {
|
||||
require.FailNowf(t, "chat should not fail", "last_error=%q", chatResult.LastError.String)
|
||||
}
|
||||
|
||||
// The LLM should have been called at least once.
|
||||
require.Greater(t, callCount.Load(), int32(0),
|
||||
"LLM should be called even when MCP token refresh fails")
|
||||
|
||||
// The original token should be unchanged in the database.
|
||||
dbToken, err := db.GetMCPServerUserToken(ctx, database.GetMCPServerUserTokenParams{
|
||||
MCPServerConfigID: mcpConfig.ID,
|
||||
UserID: user.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "old-expired-token", dbToken.AccessToken,
|
||||
"original token should be preserved when refresh fails")
|
||||
}
|
||||
|
||||
func TestChatTemplateAllowlistEnforcement(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package mcpclient
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
@@ -17,6 +18,7 @@ import (
|
||||
"github.com/mark3labs/mcp-go/client"
|
||||
"github.com/mark3labs/mcp-go/client/transport"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"golang.org/x/oauth2"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
@@ -635,3 +637,81 @@ func convertCallResult(
|
||||
}
|
||||
return fantasy.NewTextResponse("")
|
||||
}
|
||||
|
||||
// RefreshResult contains the outcome of an OAuth2 token refresh
|
||||
// attempt.
|
||||
type RefreshResult struct {
|
||||
// AccessToken is the new (or unchanged) access token.
|
||||
AccessToken string
|
||||
// RefreshToken is the new (or preserved original) refresh
|
||||
// token. Providers that don't rotate refresh tokens return
|
||||
// an empty value; in that case the original is kept.
|
||||
RefreshToken string
|
||||
// TokenType is the token type (usually "Bearer").
|
||||
TokenType string
|
||||
// Expiry is the new token expiry. Zero value means no expiry
|
||||
// was provided by the provider.
|
||||
Expiry time.Time
|
||||
// Refreshed is true when the access token actually changed,
|
||||
// meaning a refresh occurred. When false the token was still
|
||||
// valid and no network call was made.
|
||||
Refreshed bool
|
||||
}
|
||||
|
||||
// RefreshOAuth2Token checks whether the given MCP user token is
|
||||
// expired (or within 10 seconds of expiry) and refreshes it using
|
||||
// the OAuth2 credentials from the server config. If the token is
|
||||
// still valid, no network call is made and Refreshed is false.
|
||||
//
|
||||
// The caller is responsible for persisting the result when
|
||||
// Refreshed is true.
|
||||
func RefreshOAuth2Token(
|
||||
ctx context.Context,
|
||||
cfg database.MCPServerConfig,
|
||||
tok database.MCPServerUserToken,
|
||||
) (RefreshResult, error) {
|
||||
oauth2Cfg := &oauth2.Config{
|
||||
ClientID: cfg.OAuth2ClientID,
|
||||
ClientSecret: cfg.OAuth2ClientSecret,
|
||||
Endpoint: oauth2.Endpoint{
|
||||
TokenURL: cfg.OAuth2TokenURL,
|
||||
},
|
||||
}
|
||||
|
||||
oldToken := &oauth2.Token{
|
||||
AccessToken: tok.AccessToken,
|
||||
RefreshToken: tok.RefreshToken,
|
||||
TokenType: tok.TokenType,
|
||||
}
|
||||
if tok.Expiry.Valid {
|
||||
oldToken.Expiry = tok.Expiry.Time
|
||||
}
|
||||
|
||||
// Cap the refresh HTTP call so a stalled token endpoint
|
||||
// cannot block the entire MCP connection phase. The timeout
|
||||
// matches connectTimeout used for MCP server connections.
|
||||
refreshCtx, cancel := context.WithTimeout(ctx, connectTimeout)
|
||||
defer cancel()
|
||||
|
||||
// TokenSource automatically refreshes expired tokens. It
|
||||
// uses a 10-second expiry window, so tokens about to expire
|
||||
// are also refreshed proactively.
|
||||
newToken, err := oauth2Cfg.TokenSource(refreshCtx, oldToken).Token()
|
||||
if err != nil {
|
||||
return RefreshResult{}, xerrors.Errorf("refresh oauth2 token: %w", err)
|
||||
}
|
||||
|
||||
refreshed := newToken.AccessToken != tok.AccessToken
|
||||
|
||||
// Preserve the old refresh token when the provider doesn't
|
||||
// rotate (returns empty).
|
||||
refreshToken := cmp.Or(newToken.RefreshToken, tok.RefreshToken)
|
||||
|
||||
return RefreshResult{
|
||||
AccessToken: newToken.AccessToken,
|
||||
RefreshToken: refreshToken,
|
||||
TokenType: newToken.TokenType,
|
||||
Expiry: newToken.Expiry,
|
||||
Refreshed: refreshed,
|
||||
}, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user