github--github-mcp-server/pkg/github/ui_capability_test.go

package github

import (
	"context"
	"testing"

	ghcontext "github.com/github/github-mcp-server/pkg/context"
	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func createMCPRequestWithCapabilities(t *testing.T, caps *mcp.ClientCapabilities) mcp.CallToolRequest {
	t.Helper()
	srv := mcp.NewServer(&mcp.Implementation{Name: "test"}, nil)
	st, _ := mcp.NewInMemoryTransports()
	session, err := srv.Connect(context.Background(), st, &mcp.ServerSessionOptions{
		State: &mcp.ServerSessionState{
			InitializeParams: &mcp.InitializeParams{
				ClientInfo:   &mcp.Implementation{Name: "test-client"},
				Capabilities: caps,
			},
		},
	})
	require.NoError(t, err)
	t.Cleanup(func() { _ = session.Close() })
	return mcp.CallToolRequest{Session: session}
}

func Test_clientSupportsUI(t *testing.T) {
	t.Parallel()
	ctx := context.Background()

	t.Run("client with UI extension", func(t *testing.T) {
		caps := &mcp.ClientCapabilities{}
		caps.AddExtension("io.modelcontextprotocol/ui", map[string]any{
			"mimeTypes": []string{"text/html;profile=mcp-app"},
		})
		req := createMCPRequestWithCapabilities(t, caps)
		assert.True(t, clientSupportsUI(ctx, &req))
	})

	t.Run("client without UI extension", func(t *testing.T) {
		req := createMCPRequestWithCapabilities(t, &mcp.ClientCapabilities{})
		assert.False(t, clientSupportsUI(ctx, &req))
	})

	t.Run("client with nil capabilities", func(t *testing.T) {
		req := createMCPRequestWithCapabilities(t, nil)
		assert.False(t, clientSupportsUI(ctx, &req))
	})

	t.Run("nil request", func(t *testing.T) {
		assert.False(t, clientSupportsUI(ctx, nil))
	})

	t.Run("nil session", func(t *testing.T) {
		req := createMCPRequest(nil)
		assert.False(t, clientSupportsUI(ctx, &req))
	})
}

func Test_clientSupportsUI_fromContext(t *testing.T) {
	t.Parallel()

	t.Run("UI supported in context", func(t *testing.T) {
		ctx := ghcontext.WithUISupport(context.Background(), true)
		assert.True(t, clientSupportsUI(ctx, nil))
	})

	t.Run("UI not supported in context", func(t *testing.T) {
		ctx := ghcontext.WithUISupport(context.Background(), false)
		assert.False(t, clientSupportsUI(ctx, nil))
	})

	t.Run("context takes precedence over session", func(t *testing.T) {
		ctx := ghcontext.WithUISupport(context.Background(), false)
		caps := &mcp.ClientCapabilities{}
		caps.AddExtension("io.modelcontextprotocol/ui", map[string]any{})
		req := createMCPRequestWithCapabilities(t, caps)
		assert.False(t, clientSupportsUI(ctx, &req))
	})

	t.Run("no context or session", func(t *testing.T) {
		assert.False(t, clientSupportsUI(context.Background(), nil))
	})
}

func Test_shouldDeferToForm_featureFlags(t *testing.T) {
	t.Parallel()

	ctx := ghcontext.WithUISupport(context.Background(), true)
	args := map[string]any{"owner": "octocat"}
	formParams := map[string]struct{}{"owner": {}}

	tests := []struct {
		name         string
		enabledFlags []string
		want         bool
	}{
		{
			name:         "MCP Apps enabled defers to form",
			enabledFlags: []string{MCPAppsFeatureFlag},
			want:         true,
		},
		{
			name: "form deferral disabled executes directly",
			enabledFlags: []string{
				MCPAppsFeatureFlag,
				MCPAppsDisableFormDeferralFeatureFlag,
			},
			want: false,
		},
		{
			name:         "form deferral opt-out does not enable MCP Apps",
			enabledFlags: []string{MCPAppsDisableFormDeferralFeatureFlag},
			want:         false,
		},
		{
			name: "MCP Apps disabled executes directly",
			want: false,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			deps := BaseDeps{featureChecker: featureCheckerFor(tc.enabledFlags...)}
			assert.Equal(t, tc.want, shouldDeferToForm(ctx, deps, nil, args, formParams))
		})
	}
}