arp/server.go

package main

import (
	"context"
	"embed"
	"log"
	"net/http"
	"os"

	"github.com/99designs/gqlgen/graphql/handler"
	"github.com/99designs/gqlgen/graphql/handler/extension"
	"github.com/99designs/gqlgen/graphql/handler/lru"
	"github.com/99designs/gqlgen/graphql/handler/transport"
	"github.com/99designs/gqlgen/graphql/playground"
	"github.com/vektah/gqlparser/v2"
	"github.com/vektah/gqlparser/v2/ast"
	"gogs.dmsc.dev/arp/auth"
	"gogs.dmsc.dev/arp/graph"
	"gogs.dmsc.dev/arp/mcp"
	"gogs.dmsc.dev/arp/models"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

//go:embed graph/schema.graphqls
var schemaFS embed.FS

const defaultPort = "8080"

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = defaultPort
	}

	// Initialize database
	db, err := gorm.Open(sqlite.Open("arp.db"), &gorm.Config{})
	if err != nil {
		log.Fatal("failed to connect database:", err)
	}

	// Run auto-migration for all models
	err = db.AutoMigrate(&models.User{}, &models.Role{}, &models.Permission{}, &models.Service{}, &models.Task{}, &models.TaskStatus{}, &models.Message{}, &models.Note{})
	if err != nil {
		log.Fatal("failed to migrate database:", err)
	}

	// Create resolver with DB instance using NewResolver for pub/sub support
	resolver := graph.NewResolver(db)

	srv := handler.New(graph.NewExecutableSchema(graph.Config{Resolvers: resolver}))

	srv.AddTransport(transport.Options{})
	srv.AddTransport(transport.GET{})
	srv.AddTransport(transport.POST{})
	srv.AddTransport(&transport.Websocket{
		// Authenticate WebSocket connections
		InitFunc: func(ctx context.Context, initPayload transport.InitPayload) (context.Context, *transport.InitPayload, error) {
			// Get token from connection params
			token, ok := initPayload["Authorization"].(string)
			if !ok {
				return ctx, nil, nil
			}

			// Validate token and add user to context
			claims, err := auth.ValidateToken(token)
			if err != nil {
				return ctx, nil, nil
			}

			// Convert Claims to UserContext
			userCtx := &auth.UserContext{
				ID:          claims.UserID,
				Email:       claims.Email,
				Roles:       claims.Roles,
				Permissions: claims.Permissions,
			}

			ctx = auth.WithUser(ctx, userCtx)
			return ctx, nil, nil
		},
	})

	srv.SetQueryCache(lru.New[*ast.QueryDocument](1000))

	srv.Use(extension.Introspection{})
	srv.Use(extension.AutomaticPersistedQuery{
		Cache: lru.New[string](100),
	})

	// Load GraphQL schema for MCP server (from embedded FS)
	schemaStr, err := schemaFS.ReadFile("graph/schema.graphqls")
	if err != nil {
		log.Fatal("failed to read embedded schema file:", err)
	}
	schema := gqlparser.MustLoadSchema(&ast.Source{
		Name:  "schema.graphqls",
		Input: string(schemaStr),
	})

	// Create MCP server
	mcpServer := mcp.NewServer(resolver, schema)

	http.Handle("/", playground.Handler("GraphQL playground", "/query"))
	http.Handle("/query", auth.AuthMiddleware(srv))
	http.Handle("/mcp", auth.AuthMiddleware(mcpServer))
	http.Handle("/message", auth.AuthMiddleware(http.HandlerFunc(mcpServer.HandleMessage)))

	log.Printf("connect to http://localhost:%s/ for GraphQL playground", port)
	log.Printf("MCP server available at http://localhost:%s/mcp", port)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}