gitea-mcp/operation/sdk_integration_test.go

package operation

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	mcpContext "gitea.com/gitea/gitea-mcp/pkg/context"
	"gitea.com/gitea/gitea-mcp/pkg/flag"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

const testServerVersion = "test-version"

func exposeAllTools(t *testing.T) {
	t.Helper()
	originalReadOnly := flag.ReadOnly
	originalAllowedTools := flag.AllowedTools
	originalAllowedScopes := flag.AllowedScopes
	originalVersion := flag.Version
	t.Cleanup(func() {
		flag.ReadOnly = originalReadOnly
		flag.AllowedTools = originalAllowedTools
		flag.AllowedScopes = originalAllowedScopes
		flag.Version = originalVersion
	})
	flag.ReadOnly = false
	flag.AllowedTools = nil
	flag.AllowedScopes = nil
	flag.Version = testServerVersion
}

func assertVersionToolResult(t *testing.T, result *mcp.CallToolResult) {
	t.Helper()
	if len(result.Content) != 1 {
		t.Fatalf("version tool content count = %d, want 1", len(result.Content))
	}
	content, ok := result.Content[0].(*mcp.TextContent)
	if !ok {
		t.Fatalf("version tool content type = %T, want *mcp.TextContent", result.Content[0])
	}
	if !strings.Contains(content.Text, testServerVersion) {
		t.Errorf("version tool result = %q, want it to contain %q", content.Text, testServerVersion)
	}
}

func listAndCallVersion(ctx context.Context, t *testing.T, session *mcp.ClientSession) {
	t.Helper()
	result, err := session.ListTools(ctx, nil)
	if err != nil {
		t.Fatalf("ListTools() error = %v", err)
	}
	if len(result.Tools) != 54 {
		t.Fatalf("ListTools() count = %d, want 54", len(result.Tools))
	}
	callResult, err := session.CallTool(ctx, &mcp.CallToolParams{
		Name: "get_gitea_mcp_server_version",
	})
	if err != nil {
		t.Fatalf("CallTool() error = %v", err)
	}
	assertVersionToolResult(t, callResult)
}

func TestOfficialSDKInMemory(t *testing.T) {
	exposeAllTools(t)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	serverTransport, clientTransport := mcp.NewInMemoryTransports()
	server := newMCPServer(testServerVersion)
	RegisterTool(server)
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- server.Run(ctx, serverTransport)
	}()

	client := mcp.NewClient(&mcp.Implementation{Name: "gitea-mcp-test", Version: "1"}, nil)
	session, err := client.Connect(ctx, clientTransport, nil)
	if err != nil {
		t.Fatalf("Connect() error = %v", err)
	}
	if got := session.InitializeResult().ProtocolVersion; got != "2026-07-28" {
		t.Errorf("protocol version = %q, want %q", got, "2026-07-28")
	}
	listAndCallVersion(ctx, t, session)
	if err := session.Close(); err != nil {
		t.Fatalf("Close() error = %v", err)
	}
	select {
	case err := <-serverDone:
		if err != nil && !errors.Is(err, context.Canceled) {
			t.Fatalf("server Run() error = %v", err)
		}
	case <-ctx.Done():
		t.Fatal("server did not stop after the client session closed")
	}
}

func TestStreamableHTTPStateful(t *testing.T) {
	exposeAllTools(t)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	server := newMCPServer(testServerVersion)
	RegisterTool(server)
	httpTestServer := httptest.NewServer(newHTTPServer("", server).Handler)
	defer httpTestServer.Close()

	client := mcp.NewClient(&mcp.Implementation{Name: "gitea-mcp-http-test", Version: "1"}, nil)
	session, err := client.Connect(ctx, &mcp.StreamableClientTransport{
		Endpoint:             httpTestServer.URL + "/mcp",
		HTTPClient:           httpTestServer.Client(),
		DisableStandaloneSSE: true,
		MaxRetries:           -1,
	}, nil)
	if err != nil {
		t.Fatalf("Connect() error = %v", err)
	}
	defer session.Close()
	if got := session.InitializeResult().ProtocolVersion; got != "2025-11-25" {
		t.Errorf("protocol version = %q, want %q", got, "2025-11-25")
	}
	listAndCallVersion(ctx, t, session)

	response, err := httpTestServer.Client().Get(httpTestServer.URL + "/not-mcp")
	if err != nil {
		t.Fatalf("GET outside /mcp error = %v", err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusNotFound {
		t.Errorf("GET outside /mcp status = %d, want %d", response.StatusCode, http.StatusNotFound)
	}
}

func TestStreamableHTTPAllowsLegacyLargeBodies(t *testing.T) {
	server := newMCPServer(testServerVersion)
	httpTestServer := httptest.NewServer(newHTTPServer("", server).Handler)
	defer httpTestServer.Close()

	body := strings.NewReader(strings.Repeat(" ", mcp.DefaultMaxRequestBodyBytes+1))
	request, err := http.NewRequest(http.MethodPost, httpTestServer.URL+"/mcp", body)
	if err != nil {
		t.Fatalf("NewRequest() error = %v", err)
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Accept", "application/json, text/event-stream")
	response, err := httpTestServer.Client().Do(request)
	if err != nil {
		t.Fatalf("POST large body error = %v", err)
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusRequestEntityTooLarge {
		t.Errorf("POST large body status = %d; PR 1 must preserve the previous unlimited body behavior", response.StatusCode)
	}
}

type authorizationTransport struct {
	base  http.RoundTripper
	mu    sync.RWMutex
	value string
}

func (t *authorizationTransport) set(value string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.value = value
}

func (t *authorizationTransport) RoundTrip(request *http.Request) (*http.Response, error) {
	clone := request.Clone(request.Context())
	clone.Header = request.Header.Clone()
	t.mu.RLock()
	value := t.value
	t.mu.RUnlock()
	if value != "" {
		clone.Header.Set("Authorization", value)
	}
	return t.base.RoundTrip(clone)
}

func authContextValue(ctx context.Context, session *mcp.ClientSession) (string, error) {
	result, err := session.CallTool(ctx, &mcp.CallToolParams{Name: "test_auth_context"})
	if err != nil {
		return "", err
	}
	if len(result.Content) != 1 {
		return "", fmt.Errorf("content count = %d, want 1", len(result.Content))
	}
	content, ok := result.Content[0].(*mcp.TextContent)
	if !ok {
		return "", fmt.Errorf("content type = %T, want *mcp.TextContent", result.Content[0])
	}
	return content.Text, nil
}

func TestHTTPAuthPerRequest(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	server := newMCPServer(testServerVersion)
	server.AddTool(
		&mcp.Tool{
			Name:        "test_auth_context",
			Description: "Return the request-scoped authentication token.",
			InputSchema: map[string]any{"type": "object", "properties": map[string]any{}},
		},
		func(ctx context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
			token, _ := ctx.Value(mcpContext.TokenContextKey).(string)
			return &mcp.CallToolResult{
				Content: []mcp.Content{&mcp.TextContent{Text: token}},
			}, nil
		},
	)
	httpTestServer := httptest.NewServer(newHTTPServer("", server).Handler)
	defer httpTestServer.Close()

	baseTransport := httpTestServer.Client().Transport
	auth := &authorizationTransport{base: baseTransport}
	auth.set("Bearer first-token")
	baseClient := &http.Client{Transport: auth}
	client := mcp.NewClient(&mcp.Implementation{Name: "gitea-mcp-auth-test", Version: "1"}, nil)
	session, err := client.Connect(ctx, &mcp.StreamableClientTransport{
		Endpoint:             httpTestServer.URL + "/mcp",
		HTTPClient:           baseClient,
		DisableStandaloneSSE: true,
		MaxRetries:           -1,
	}, nil)
	if err != nil {
		t.Fatalf("Connect() error = %v", err)
	}
	defer session.Close()

	for _, test := range []struct {
		header string
		want   string
	}{
		{header: "Bearer first-token", want: "first-token"},
		{header: "token second-token", want: "second-token"},
		{header: "Basic ignored", want: ""},
	} {
		auth.set(test.header)
		token, err := authContextValue(ctx, session)
		if err != nil {
			t.Fatalf("CallTool() with %q error = %v", test.header, err)
		}
		if token != test.want {
			t.Errorf("CallTool() token = %q, want %q", token, test.want)
		}
	}

	type authenticatedSession struct {
		session *mcp.ClientSession
		want    string
	}
	concurrentSessions := make([]authenticatedSession, 0, 2)
	for index, token := range []string{"parallel-one", "parallel-two"} {
		transport := &authorizationTransport{base: baseTransport}
		transport.set("Bearer " + token)
		httpClient := &http.Client{Transport: transport}
		parallelClient := mcp.NewClient(&mcp.Implementation{
			Name:    fmt.Sprintf("gitea-mcp-auth-parallel-%d", index),
			Version: "1",
		}, nil)
		parallelSession, err := parallelClient.Connect(ctx, &mcp.StreamableClientTransport{
			Endpoint:             httpTestServer.URL + "/mcp",
			HTTPClient:           httpClient,
			DisableStandaloneSSE: true,
			MaxRetries:           -1,
		}, nil)
		if err != nil {
			t.Fatalf("parallel Connect() error = %v", err)
		}
		defer parallelSession.Close()
		concurrentSessions = append(concurrentSessions, authenticatedSession{session: parallelSession, want: token})
	}

	var waitGroup sync.WaitGroup
	errorsCh := make(chan error, 20)
	for _, authenticated := range concurrentSessions {
		for range 10 {
			waitGroup.Go(func() {
				got, err := authContextValue(ctx, authenticated.session)
				if err != nil {
					errorsCh <- err
					return
				}
				if got != authenticated.want {
					errorsCh <- fmt.Errorf("parallel token = %q, want %q", got, authenticated.want)
				}
			})
		}
	}
	waitGroup.Wait()
	close(errorsCh)
	for err := range errorsCh {
		t.Error(err)
	}
}

func TestStdioCommandTransport(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping subprocess build in short mode")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	binary := filepath.Join(t.TempDir(), "gitea-mcp")
	build := exec.CommandContext(ctx, "go", "build", "-o", binary, "..")
	build.Env = os.Environ()
	if output, err := build.CombinedOutput(); err != nil {
		t.Fatalf("build stdio test binary: %v\n%s", err, output)
	}

	client := mcp.NewClient(&mcp.Implementation{Name: "gitea-mcp-stdio-test", Version: "1"}, nil)
	command := exec.CommandContext(ctx, binary, "--transport", "stdio")
	session, err := client.Connect(ctx, &mcp.CommandTransport{
		Command:           command,
		TerminateDuration: 2 * time.Second,
	}, nil)
	if err != nil {
		t.Fatalf("Connect() error = %v", err)
	}
	defer session.Close()
	if got := session.InitializeResult().ProtocolVersion; got != "2026-07-28" {
		t.Errorf("protocol version = %q, want %q", got, "2026-07-28")
	}
	result, err := session.ListTools(ctx, nil)
	if err != nil {
		t.Fatalf("ListTools() error = %v", err)
	}
	if len(result.Tools) != 54 {
		t.Errorf("ListTools() count = %d, want 54", len(result.Tools))
	}
	callResult, err := session.CallTool(ctx, &mcp.CallToolParams{Name: "get_gitea_mcp_server_version"})
	if err != nil {
		t.Fatalf("CallTool() error = %v", err)
	}
	content, ok := callResult.Content[0].(*mcp.TextContent)
	if !ok {
		t.Fatalf("CallTool() content type = %T, want *mcp.TextContent", callResult.Content[0])
	}
	if !strings.Contains(content.Text, "Gitea MCP Server version:") {
		t.Errorf("CallTool() result = %q, want server version", content.Text)
	}
}

func TestNewHTTPServerAddress(t *testing.T) {
	server := newHTTPServer(":12345", newMCPServer(testServerVersion))
	if server.Addr != ":12345" {
		t.Errorf("server address = %q, want %q", server.Addr, ":12345")
	}
	if server.Handler == nil {
		t.Error("server handler is nil")
	}
}

func TestHTTPServerGracefulShutdown(t *testing.T) {
	server := newHTTPServer("127.0.0.1:0", newMCPServer(testServerVersion))
	listener, err := net.Listen("tcp", server.Addr)
	if err != nil {
		t.Fatalf("Listen() error = %v", err)
	}
	serveDone := make(chan error, 1)
	go func() {
		serveDone <- server.Serve(listener)
	}()

	shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := server.Shutdown(shutdownCtx); err != nil {
		t.Fatalf("Shutdown() error = %v", err)
	}
	select {
	case err := <-serveDone:
		if !errors.Is(err, http.ErrServerClosed) {
			t.Errorf("Serve() error = %v, want http.ErrServerClosed", err)
		}
	case <-shutdownCtx.Done():
		t.Fatal("server did not stop after Shutdown()")
	}
}