munin/internal/mcpserver/concurrency_test.go

package mcpserver

import (
	"context"
	"fmt"
	"sync"
	"testing"

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

	"munin/internal/auth"
	"munin/internal/store"
	"munin/internal/web"
	"munin/internal/workspace"
)

// connect öffnet eine MCP-Sitzung, über die man mehrere Aufrufe schicken kann.
//
// Anders als call meldet es Fehler zurück, statt t.Fatalf aufzurufen — t.Fatalf
// darf nur aus der Goroutine des Tests selbst kommen, nicht aus den
// nebenläufigen Agenten unten.
func connect(h *Handler, ws *workspace.Workspace, project store.Project, name string) (*mcp.ClientSession, func(), error) {
	sess := web.Session{
		Subject:     "key:" + ws.ID + ":" + name,
		WorkspaceID: ws.ID,
		Key:         store.APIKey{ID: name, Name: name, ProjectID: project.ID, Role: string(auth.RoleWrite)},
		Project:     project,
		Role:        auth.RoleWrite,
	}
	ctx := web.WithSessionForTest(context.Background(), sess)
	ctx = web.WithWorkspaceForTest(ctx, ws)

	clientTransport, serverTransport := mcp.NewInMemoryTransports()
	serverSession, err := h.serverFor(project).Connect(ctx, serverTransport, nil)
	if err != nil {
		return nil, nil, err
	}
	client := mcp.NewClient(&mcp.Implementation{Name: name, Version: "1"}, nil)
	clientSession, err := client.Connect(ctx, clientTransport, nil)
	if err != nil {
		serverSession.Close()
		return nil, nil, err
	}
	closeIt := func() {
		clientSession.Close()
		serverSession.Close()
	}
	return clientSession, closeIt, nil
}

// TestConcurrentWritesKeepSearchIndexCurrent prüft, dass der Suchindex
// nach gleichzeitigen Schreibvorgängen auf denselben Eintrag den Stand der
// Datenbank kennt.
//
// Die Gefahr: remember schreibt erst in die Datenbank und aktualisiert danach
// den Index — zwei getrennte Schritte. Überholen sich zwei Agenten zwischen
// diesen Schritten, landet Fassung 3 in der Datenbank, aber Fassung 2 im Index.
// recall fände den Eintrag dann bis zum Neustart über veraltete Wörter.
//
// Jeder Schreibvorgang bekommt deshalb ein eigenes, eindeutiges Wort. Am Ende
// muss das Wort der letzten Fassung in der Datenbank auch im Index zu finden
// sein.
func TestConcurrentWritesKeepSearchIndexCurrent(t *testing.T) {
	h, ws, project := setupWithBudget(t, 100000)
	ctx := context.Background()

	// Den Index vorher anlegen. Ohne ihn wäre Search.Update ein No-op und der
	// Test würde nichts prüfen (siehe Registry.Update).
	idx, err := ws.Search.For(ctx, project.ID)
	if err != nil {
		t.Fatal(err)
	}

	const rounds, agents, perAgent = 20, 8, 10

	for round := 0; round < rounds; round++ {
		slug := fmt.Sprintf("umkaempft-%02d", round)

		var wg sync.WaitGroup
		var mu sync.Mutex
		var errs []error

		for a := 0; a < agents; a++ {
			wg.Add(1)
			go func(a int) {
				defer wg.Done()
				report := func(err error) {
					mu.Lock()
					errs = append(errs, err)
					mu.Unlock()
				}

				client, closeIt, err := connect(h, ws, project, fmt.Sprintf("agent-%d", a))
				if err != nil {
					report(err)
					return
				}
				defer closeIt()

				for i := 0; i < perAgent; i++ {
					// Nur Buchstaben, damit der Tokenizer das Wort nicht zerlegt.
					word := fmt.Sprintf("marke%sx%s", letters(round*1000+a*100+i), "ende")
					res, err := client.CallTool(ctx, &mcp.CallToolParams{
						Name: "remember",
						// force: Jede Runde legt unter neuem Bezeichner mit
						// demselben Titel an. Die Dublettenprüfung würde das
						// zu Recht anhalten — hier geht es aber ums
						// gleichzeitige Schreiben, nicht um Dubletten.
						Arguments: map[string]any{
							"slug": slug, "title": "Umkämpft", "content": "Stand " + word, "force": true,
						},
					})
					if err != nil {
						report(err)
						return
					}
					if res.IsError {
						report(fmt.Errorf("remember: %s", text(res)))
						return
					}
				}
			}(a)
		}
		wg.Wait()

		if len(errs) > 0 {
			t.Fatalf("Runde %d: %d Fehler, erster: %v", round, len(errs), errs[0])
		}

		e, err := ws.Store.EntryBySlug(ctx, project.ID, slug)
		if err != nil {
			t.Fatal(err)
		}
		if e.Version != agents*perAgent {
			t.Errorf("Runde %d: Fassung %d, erwartet %d — ein Schreibvorgang ging verloren",
				round, e.Version, agents*perAgent)
		}

		// Das Wort der letzten Fassung steht nach "Stand " im Inhalt.
		lastWord := e.Content[len("Stand "):]
		found := false
		for _, r := range idx.Search(lastWord, 5) {
			if r.DocID == slug {
				found = true
			}
		}
		if !found {
			t.Errorf("Runde %d: Datenbank hat Fassung %d mit %q, der Suchindex kennt das Wort nicht — "+
				"er hält einen älteren Stand", round, e.Version, lastWord)
		}
	}
}

// letters schreibt eine Zahl mit den Buchstaben a–j statt der Ziffern 0–9.
// So entsteht für jeden Schreibvorgang ein eindeutiges Wort ohne Ziffern.
func letters(n int) string {
	s := fmt.Sprintf("%06d", n)
	b := []byte(s)
	for i := range b {
		b[i] = 'a' + (b[i] - '0')
	}
	return string(b)
}