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.
(t *testing.T)
| 8116 | // connecting. The refreshed token is persisted to the database and |
| 8117 | // the MCP tool call succeeds. |
| 8118 | func TestMCPServerOAuth2TokenRefresh(t *testing.T) { |
| 8119 | t.Parallel() |
| 8120 | |
| 8121 | db, ps := dbtestutil.NewDB(t) |
| 8122 | ctx := testutil.Context(t, testutil.WaitLong) |
| 8123 | |
| 8124 | // The "fresh" token that the mock OAuth2 server returns after |
| 8125 | // a successful refresh_token grant. |
| 8126 | freshAccessToken := "fresh-access-token-" + uuid.New().String() |
| 8127 | |
| 8128 | // Mock OAuth2 token endpoint that exchanges a refresh token |
| 8129 | // for a new access token. |
| 8130 | var refreshCalled atomic.Int32 |
| 8131 | tokenSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 8132 | refreshCalled.Add(1) |
| 8133 | |
| 8134 | if r.Method != http.MethodPost { |
| 8135 | w.WriteHeader(http.StatusMethodNotAllowed) |
| 8136 | return |
| 8137 | } |
| 8138 | |
| 8139 | grantType := r.FormValue("grant_type") |
| 8140 | if grantType != "refresh_token" { |
| 8141 | w.Header().Set("Content-Type", "application/json") |
| 8142 | w.WriteHeader(http.StatusBadRequest) |
| 8143 | _, _ = w.Write([]byte(`{"error":"unsupported_grant_type"}`)) |
| 8144 | return |
| 8145 | } |
| 8146 | |
| 8147 | w.Header().Set("Content-Type", "application/json") |
| 8148 | _, _ = fmt.Fprintf(w, `{"access_token":%q,"token_type":"Bearer","expires_in":3600,"refresh_token":"rotated-refresh-token"}`, freshAccessToken) |
| 8149 | })) |
| 8150 | t.Cleanup(tokenSrv.Close) |
| 8151 | |
| 8152 | // Start a real MCP server with an auth middleware that only |
| 8153 | // accepts the fresh access token. An expired token (or any |
| 8154 | // other value) gets a 401. |
| 8155 | mcpSrv := mcpserver.NewMCPServer("authed-mcp", "1.0.0") |
| 8156 | mcpSrv.AddTools(mcpserver.ServerTool{ |
| 8157 | Tool: mcpgo.NewTool("echo", |
| 8158 | mcpgo.WithDescription("Echoes the input"), |
| 8159 | mcpgo.WithString("input", |
| 8160 | mcpgo.Description("The input string"), |
| 8161 | mcpgo.Required(), |
| 8162 | ), |
| 8163 | ), |
| 8164 | Handler: func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { |
| 8165 | input, _ := req.GetArguments()["input"].(string) |
| 8166 | return mcpgo.NewToolResultText("echo: " + input), nil |
| 8167 | }, |
| 8168 | }) |
| 8169 | mcpHTTP := mcpserver.NewStreamableHTTPServer(mcpSrv) |
| 8170 | // Wrap with auth check. |
| 8171 | authMux := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 8172 | auth := r.Header.Get("Authorization") |
| 8173 | if auth != "Bearer "+freshAccessToken { |
| 8174 | w.WriteHeader(http.StatusUnauthorized) |
| 8175 | _, _ = w.Write([]byte(`{"error":"invalid_token","error_description":"The access token is invalid or expired"}`)) |
nothing calls this directly
no test coverage detected