@@ -4,58 +4,84 @@ import (
44 "context"
55 "testing"
66
7+ ghcontext "github.com/github/github-mcp-server/pkg/context"
78 "github.com/modelcontextprotocol/go-sdk/mcp"
89 "github.com/stretchr/testify/assert"
10+ "github.com/stretchr/testify/require"
911)
1012
13+ func createMCPRequestWithCapabilities (t * testing.T , caps * mcp.ClientCapabilities ) mcp.CallToolRequest {
14+ t .Helper ()
15+ srv := mcp .NewServer (& mcp.Implementation {Name : "test" }, nil )
16+ st , _ := mcp .NewInMemoryTransports ()
17+ session , err := srv .Connect (context .Background (), st , & mcp.ServerSessionOptions {
18+ State : & mcp.ServerSessionState {
19+ InitializeParams : & mcp.InitializeParams {
20+ ClientInfo : & mcp.Implementation {Name : "test-client" },
21+ Capabilities : caps ,
22+ },
23+ },
24+ })
25+ require .NoError (t , err )
26+ t .Cleanup (func () { _ = session .Close () })
27+ return mcp.CallToolRequest {Session : session }
28+ }
29+
1130func Test_clientSupportsUI (t * testing.T ) {
1231 t .Parallel ()
32+ ctx := context .Background ()
1333
14- tests := []struct {
15- name string
16- clientName string
17- want bool
18- }{
19- {name : "VS Code Insiders" , clientName : "Visual Studio Code - Insiders" , want : true },
20- {name : "VS Code Stable" , clientName : "Visual Studio Code" , want : true },
21- {name : "unknown client" , clientName : "some-other-client" , want : false },
22- {name : "empty client name" , clientName : "" , want : false },
23- }
24-
25- for _ , tt := range tests {
26- t .Run (tt .name , func (t * testing.T ) {
27- req := createMCPRequestWithSession (t , tt .clientName , nil )
28- assert .Equal (t , tt .want , clientSupportsUI (& req ))
34+ t .Run ("client with UI extension" , func (t * testing.T ) {
35+ caps := & mcp.ClientCapabilities {}
36+ caps .AddExtension ("io.modelcontextprotocol/ui" , map [string ]any {
37+ "mimeTypes" : []string {"text/html;profile=mcp-app" },
2938 })
30- }
39+ req := createMCPRequestWithCapabilities (t , caps )
40+ assert .True (t , clientSupportsUI (ctx , & req ))
41+ })
42+
43+ t .Run ("client without UI extension" , func (t * testing.T ) {
44+ req := createMCPRequestWithCapabilities (t , & mcp.ClientCapabilities {})
45+ assert .False (t , clientSupportsUI (ctx , & req ))
46+ })
47+
48+ t .Run ("client with nil capabilities" , func (t * testing.T ) {
49+ req := createMCPRequestWithCapabilities (t , nil )
50+ assert .False (t , clientSupportsUI (ctx , & req ))
51+ })
3152
3253 t .Run ("nil request" , func (t * testing.T ) {
33- assert .False (t , clientSupportsUI (nil ))
54+ assert .False (t , clientSupportsUI (ctx , nil ))
3455 })
3556
3657 t .Run ("nil session" , func (t * testing.T ) {
3758 req := createMCPRequest (nil )
38- assert .False (t , clientSupportsUI (& req ))
59+ assert .False (t , clientSupportsUI (ctx , & req ))
3960 })
4061}
4162
42- func Test_clientSupportsUI_nilClientInfo (t * testing.T ) {
63+ func Test_clientSupportsUI_fromContext (t * testing.T ) {
4364 t .Parallel ()
4465
45- srv := mcp . NewServer ( & mcp. Implementation { Name : "test" }, nil )
46- st , _ := mcp . NewInMemoryTransports ( )
47- session , err := srv . Connect ( context . Background (), st , & mcp. ServerSessionOptions {
48- State : & mcp. ServerSessionState {
49- InitializeParams : & mcp. InitializeParams {
50- ClientInfo : nil ,
51- },
52- },
66+ t . Run ( "UI supported in context" , func ( t * testing. T ) {
67+ ctx := ghcontext . WithUISupport ( context . Background (), true )
68+ assert . True ( t , clientSupportsUI ( ctx , nil ))
69+ })
70+
71+ t . Run ( "UI not supported in context" , func ( t * testing. T ) {
72+ ctx := ghcontext . WithUISupport ( context . Background (), false )
73+ assert . False ( t , clientSupportsUI ( ctx , nil ))
5374 })
54- if err != nil {
55- t .Fatal (err )
56- }
57- t .Cleanup (func () { _ = session .Close () })
5875
59- req := mcp.CallToolRequest {Session : session }
60- assert .False (t , clientSupportsUI (& req ))
76+ t .Run ("context takes precedence over session" , func (t * testing.T ) {
77+ ctx := ghcontext .WithUISupport (context .Background (), false )
78+ caps := & mcp.ClientCapabilities {}
79+ caps .AddExtension ("io.modelcontextprotocol/ui" , map [string ]any {})
80+ req := createMCPRequestWithCapabilities (t , caps )
81+ assert .False (t , clientSupportsUI (ctx , & req ))
82+ })
83+
84+ t .Run ("no context or session" , func (t * testing.T ) {
85+ assert .False (t , clientSupportsUI (context .Background (), nil ))
86+ })
6187}
0 commit comments