petal/internal/suggestions/handlers_test.go

package suggestions

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"

	"github.com/go-chi/chi/v5"

	"gitea.parodia.dev/drwily/petal/internal/auth"
	"gitea.parodia.dev/drwily/petal/internal/db"
	"gitea.parodia.dev/drwily/petal/internal/llm"
)

// stubClient returns a canned checkpoint response; it never touches the network.
type stubClient struct {
	response string
	calls    int
	// The full prompt of the most recent call, so a test can assert which
	// sentences a chunked pass actually asked about.
	lastPrompt string
}

func (s *stubClient) Complete(_ context.Context, req llm.CompletionRequest) (string, error) {
	s.calls++
	s.lastPrompt = ""
	for _, m := range req.Messages {
		s.lastPrompt += m.Content + "\n"
	}
	return s.response, nil
}

func (s *stubClient) Stream(_ context.Context, _ llm.CompletionRequest) (<-chan string, error) {
	ch := make(chan string)
	close(ch)
	return ch, nil
}

// newTestServer wires a real DB, a seeded document, and the suggestion routes
// (both the doc-scoped and the per-suggestion mounts) onto one router.
func newTestServer(t *testing.T, client llm.LLMClient) (http.Handler, string, *Handler) {
	t.Helper()
	database, err := db.Open(filepath.Join(t.TempDir(), "test.db"))
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	t.Cleanup(func() { database.Close() })

	// Seed a document with some content to check.
	var docID string
	err = database.QueryRow(
		`INSERT INTO documents (user_id, content_text) VALUES (?, ?) RETURNING id`,
		db.LocalUserID, "I has two apple.",
	).Scan(&docID)
	if err != nil {
		t.Fatalf("seed doc: %v", err)
	}

	h := New(database, client)
	r := chi.NewRouter()
	r.Route("/docs", func(dr chi.Router) { h.RegisterDocRoutes(dr) })
	r.Mount("/suggestions", h.Routes())

	// Behind the same auth middleware main.go installs: handlers resolve the
	// caller from the request context, so a bare router would see no user.
	authed := auth.Middleware(auth.StaticResolver(db.LocalUserID))(r)
	return authed, docID, h
}

// setDocText rewrites the seeded document, standing in for the writer editing.
// The grammar checkpoint only asks the model about sentences that changed since
// it last read the document, so a test that wants a second real pass has to
// change something first — as she always has.
func setDocText(t *testing.T, h *Handler, docID, text string) {
	t.Helper()
	if _, err := h.DB.Exec(`UPDATE documents SET content_text = ? WHERE id = ?`, text, docID); err != nil {
		t.Fatalf("update doc text: %v", err)
	}
}

func do(t *testing.T, srv http.Handler, method, path, body string) *httptest.ResponseRecorder {
	t.Helper()
	var r *http.Request
	if body != "" {
		r = httptest.NewRequest(method, path, bytes.NewBufferString(body))
	} else {
		r = httptest.NewRequest(method, path, nil)
	}
	rec := httptest.NewRecorder()
	srv.ServeHTTP(rec, r)
	return rec
}

func TestCheckpointFlow(t *testing.T) {
	client := &stubClient{response: `{"suggestions":[
		{"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"},
		{"original":"two apple","replacement":"two apples","explanation":"plural noun","type":"grammar"}
	]}`}
	srv, docID, _ := newTestServer(t, client)

	// Run a checkpoint → two pending suggestions, with positions located.
	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body)
	}
	var got []db.Suggestion
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("want 2 suggestions, got %d: %+v", len(got), got)
	}
	if got[0].Original != "I has" || got[0].FromPos != 0 {
		t.Fatalf("first suggestion anchoring wrong: %+v", got[0])
	}
	if got[0].Status != db.SuggestionStatusPending {
		t.Fatalf("new suggestion should be pending: %+v", got[0])
	}

	// Listing returns the same pending set.
	rec = do(t, srv, http.MethodGet, "/docs/"+docID+"/suggestions", "")
	var listed []db.Suggestion
	_ = json.Unmarshal(rec.Body.Bytes(), &listed)
	if len(listed) != 2 {
		t.Fatalf("list pending: want 2, got %d", len(listed))
	}

	// Accept the first, dismiss the second.
	if rec = do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/accept", ""); rec.Code != http.StatusNoContent {
		t.Fatalf("accept: code=%d body=%s", rec.Code, rec.Body)
	}
	if rec = do(t, srv, http.MethodPost, "/suggestions/"+got[1].ID+"/dismiss", ""); rec.Code != http.StatusNoContent {
		t.Fatalf("dismiss: code=%d body=%s", rec.Code, rec.Body)
	}

	// Pending list is now empty (accepted/dismissed are excluded).
	rec = do(t, srv, http.MethodGet, "/docs/"+docID+"/suggestions", "")
	_ = json.Unmarshal(rec.Body.Bytes(), &listed)
	if len(listed) != 0 {
		t.Fatalf("pending after resolve: want 0, got %d", len(listed))
	}

	// Re-accepting an already-resolved suggestion 404s.
	if rec = do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/accept", ""); rec.Code != http.StatusNotFound {
		t.Fatalf("re-accept: want 404, got %d", rec.Code)
	}
}

// TestResolvedSuggestionsNotReproposed proves a checkpoint won't re-nag a
// sentence the user already accepted or dismissed: the model returns the same
// raw batch on the next pass (it has no memory), but the resolved edits are
// suppressed. This is the "accept it, then it nags again moments later" bug.
func TestResolvedSuggestionsNotReproposed(t *testing.T) {
	client := &stubClient{response: `{"suggestions":[
		{"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"},
		{"original":"two apple","replacement":"two apples","explanation":"plural noun","type":"grammar"}
	]}`}
	srv, docID, h := newTestServer(t, client)
	// Zero the checkpoint floor so the second pass runs instead of being throttled.
	h.Limit = llm.NewRateLimiter(0)

	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	var got []db.Suggestion
	_ = json.Unmarshal(rec.Body.Bytes(), &got)
	if len(got) != 2 {
		t.Fatalf("first pass: want 2, got %d", len(got))
	}

	// Accept one, dismiss the other.
	do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/accept", "")
	do(t, srv, http.MethodPost, "/suggestions/"+got[1].ID+"/dismiss", "")

	// Second checkpoint returns the identical raw batch — both must be suppressed.
	rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("second check: code=%d body=%s", rec.Code, rec.Body)
	}
	var again []db.Suggestion
	if err := json.Unmarshal(rec.Body.Bytes(), &again); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if len(again) != 0 {
		t.Fatalf("resolved suggestions should not be re-proposed, got %d: %+v", len(again), again)
	}
}

// TestFickleEditsSuppressed proves the suppressor kills the "keeps going back and
// forth on a few sentences" behavior seen live on the Missing Key doc: a later
// pass that (a) reverses an edit the user just accepted — even with the editor's
// double→single quote churn that defeats a byte-exact match — or (b) re-polishes
// the model's own accepted output. A genuinely new edit on a fresh sentence still
// survives.
func TestFickleEditsSuppressed(t *testing.T) {
	client := &stubClient{response: `{"suggestions":[
		{"original":"He left \"early,\" because of the rain.","replacement":"He left \"early,\" due to the rain.","explanation":"smoother","type":"phrasing"},
		{"original":"The cat always have a calm face.","replacement":"The cat always has a calm face.","explanation":"agreement","type":"grammar"}
	]}`}
	srv, docID, h := newTestServer(t, client)
	h.Limit = llm.NewRateLimiter(0)
	setDocText(t, h, docID, `He left "early," because of the rain. The cat always have a calm face.`)

	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	var got []db.Suggestion
	_ = json.Unmarshal(rec.Body.Bytes(), &got)
	if len(got) != 2 {
		t.Fatalf("first pass: want 2, got %d", len(got))
	}
	for _, s := range got {
		do(t, srv, http.MethodPost, "/suggestions/"+s.ID+"/accept", "")
	}

	// Both edits are now in the document, which is what re-opens those sentences
	// for a second reading.
	setDocText(t, h, docID, `He left "early," due to the rain. The cat always has a calm face.`)

	// Reversal of the first accept (note the " → ' quote churn) and a re-polish of
	// the second accept must both be dropped; only the unrelated edit survives.
	client.response = `{"suggestions":[
		{"original":"He left 'early,' due to the rain.","replacement":"He left 'early,' because of the rain.","explanation":"reverts the accepted edit","type":"phrasing"},
		{"original":"The cat always has a calm face.","replacement":"The cat always seems to have a calm face.","explanation":"re-polishes accepted output","type":"phrasing"},
		{"original":"due to the rain","replacement":"because of the rain","explanation":"sub-clause reversal of the accepted span","type":"phrasing"},
		{"original":"She drived home.","replacement":"She drove home.","explanation":"past tense","type":"grammar"}
	]}`
	rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	var again []db.Suggestion
	if err := json.Unmarshal(rec.Body.Bytes(), &again); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if len(again) != 1 || again[0].Original != "She drived home." {
		t.Fatalf("want only the new edit to survive, got %d: %+v", len(again), again)
	}
}

// TestVoicePassCoexists proves the voice pass and the grammar checkpoint own
// disjoint pending families: running one never wipes the other's flags, and
// each endpoint returns the unified pending set.
func TestVoicePassCoexists(t *testing.T) {
	// One client serves both passes; the prompt distinguishes them but the stub
	// just echoes whatever we set before each call.
	client := &switchClient{}
	srv, docID, h := newTestServer(t, client)
	// Zero the voice floor so the test can re-run the pass without waiting out
	// the real 20s interval; the family-scoping logic is what's under test here.
	h.VoiceLimit = llm.NewRateLimiter(0)

	// Grammar checkpoint first → one grammar flag.
	client.response = `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"agreement","type":"grammar"}]}`
	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body)
	}

	// Voice pass next → one voice flag (null replacement). It must NOT remove the
	// grammar flag; the response is the unified set of both.
	client.response = `{"suggestions":[{"original":"two apple","replacement":null,"explanation":"sounds more formal","type":"voice"}]}`
	rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/voice", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("voice: code=%d body=%s", rec.Code, rec.Body)
	}
	var got []db.Suggestion
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("voice response should carry both families, got %d: %+v", len(got), got)
	}
	var grammar, voice *db.Suggestion
	for i := range got {
		switch got[i].Type {
		case db.SuggestionTypeGrammar:
			grammar = &got[i]
		case db.SuggestionTypeVoice:
			voice = &got[i]
		}
	}
	if grammar == nil || voice == nil {
		t.Fatalf("want one grammar + one voice flag, got %+v", got)
	}
	if voice.Replacement != "" {
		t.Fatalf("voice flag should have empty replacement, got %q", voice.Replacement)
	}

	// A fresh voice pass that finds nothing replaces only the voice flag; the
	// grammar one survives.
	client.response = `{"suggestions":[]}`
	rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/voice", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("second voice: code=%d body=%s", rec.Code, rec.Body)
	}
	_ = json.Unmarshal(rec.Body.Bytes(), &got)
	if len(got) != 1 || got[0].Type != db.SuggestionTypeGrammar {
		t.Fatalf("after clearing voice, only the grammar flag should remain, got %+v", got)
	}
}

// TestCollocationPassCoexists proves the collocation coach is a third
// independent family: it never wipes grammar or voice pending flags, a grammar
// checkpoint never wipes its flags, and each endpoint returns the unified set.
func TestCollocationPassCoexists(t *testing.T) {
	client := &switchClient{}
	srv, docID, h := newTestServer(t, client)
	// Zero the floors so the test can re-run passes without waiting them out.
	h.Limit = llm.NewRateLimiter(0)
	h.VoiceLimit = llm.NewRateLimiter(0)
	h.CollocationLimit = llm.NewRateLimiter(0)

	// Grammar + voice first, so all three families are exercised.
	client.response = `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"agreement","type":"grammar"}]}`
	if rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", ""); rec.Code != http.StatusOK {
		t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body)
	}
	client.response = `{"suggestions":[{"original":"two apple","replacement":null,"explanation":"sounds formal","type":"voice"}]}`
	if rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/voice", ""); rec.Code != http.StatusOK {
		t.Fatalf("voice: code=%d body=%s", rec.Code, rec.Body)
	}

	// Collocation pass → a third flag (with a replacement). It must keep the
	// grammar and voice flags; the response is the unified set of all three.
	client.response = `{"suggestions":[{"original":"two apple","replacement":"two apples","explanation":"Natives usually say…","type":"collocation"}]}`
	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/collocation", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("collocation: code=%d body=%s", rec.Code, rec.Body)
	}
	var got []db.Suggestion
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode: %v", err)
	}
	byType := map[string]bool{}
	for _, s := range got {
		byType[s.Type] = true
	}
	if !byType[db.SuggestionTypeGrammar] || !byType[db.SuggestionTypeVoice] || !byType[db.SuggestionTypeCollocation] {
		t.Fatalf("collocation response should carry all three families, got %+v", got)
	}

	// A grammar checkpoint must NOT wipe the voice or collocation flags. She fixes
	// the flagged sentence, so its own grammar row goes and nothing replaces it.
	setDocText(t, h, docID, "I have two apples.")
	client.response = `{"suggestions":[]}`
	rec = do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("second check: code=%d body=%s", rec.Code, rec.Body)
	}
	_ = json.Unmarshal(rec.Body.Bytes(), &got)
	byType = map[string]bool{}
	for _, s := range got {
		byType[s.Type] = true
	}
	if byType[db.SuggestionTypeGrammar] {
		t.Fatalf("grammar flag should have cleared, got %+v", got)
	}
	if !byType[db.SuggestionTypeVoice] || !byType[db.SuggestionTypeCollocation] {
		t.Fatalf("grammar checkpoint wiped a sibling family, got %+v", got)
	}
}

// switchClient returns a response that can be swapped between calls.
type switchClient struct {
	response string
	calls    int
}

func (s *switchClient) Complete(_ context.Context, _ llm.CompletionRequest) (string, error) {
	s.calls++
	return s.response, nil
}

func (s *switchClient) Stream(_ context.Context, _ llm.CompletionRequest) (<-chan string, error) {
	ch := make(chan string)
	close(ch)
	return ch, nil
}

func TestCheckpointRateLimit(t *testing.T) {
	client := &stubClient{response: `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"x","type":"grammar"}]}`}
	srv, docID, _ := newTestServer(t, client)

	// First check runs the model.
	do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	// Immediate second check is throttled — returns the existing set without
	// hitting the model again.
	do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	if client.calls != 1 {
		t.Fatalf("rate limit should suppress the second model call, got %d calls", client.calls)
	}
}

// postMechanics submits a deterministic-findings batch (as the client's prose
// detector would) and returns the unified pending set.
func postMechanics(t *testing.T, srv http.Handler, docID, findingsJSON string) []db.Suggestion {
	t.Helper()
	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/mechanics", `{"findings":`+findingsJSON+`}`)
	if rec.Code != http.StatusOK {
		t.Fatalf("mechanics: code=%d body=%s", rec.Code, rec.Body)
	}
	var got []db.Suggestion
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode: %v", err)
	}
	return got
}

// TestMechanicsPersistsFindings proves the endpoint persists client-detected
// findings as a 'mechanics' family with the exact offsets the client supplied,
// and that a later grammar checkpoint leaves them untouched (own family).
func TestMechanicsPersistsFindings(t *testing.T) {
	client := &stubClient{response: `{"suggestions":[]}`}
	srv, docID, _ := newTestServer(t, client)

	got := postMechanics(t, srv, docID, `[
		{"from":0,"to":1,"original":"i","replacement":"I","explanation":"capitalize I"},
		{"from":6,"to":13,"original":"the the","replacement":"the","explanation":"doubled word"}
	]`)
	var mech []db.Suggestion
	for _, s := range got {
		if s.Type == db.SuggestionTypeMechanics {
			mech = append(mech, s)
		}
	}
	if len(mech) != 2 {
		t.Fatalf("want 2 mechanics findings, got %d: %+v", len(mech), got)
	}
	// The client's exact offsets are preserved verbatim.
	for _, s := range mech {
		if s.Original == "the the" && (s.FromPos != 6 || s.ToPos != 13) {
			t.Errorf("offsets not preserved: %+v", s)
		}
	}

	// A grammar checkpoint must not wipe the mechanics family.
	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body)
	}
	var after []db.Suggestion
	_ = json.Unmarshal(rec.Body.Bytes(), &after)
	count := 0
	for _, s := range after {
		if s.Type == db.SuggestionTypeMechanics {
			count++
		}
	}
	if count != 2 {
		t.Fatalf("grammar checkpoint disturbed the mechanics family: got %d, want 2", count)
	}
}

// TestMechanicsWinsSpanCollision proves that when a mechanics finding and an LLM
// grammar suggestion claim overlapping characters, the LLM card is dropped from
// the response and the exact mechanics fix owns the span.
func TestMechanicsWinsSpanCollision(t *testing.T) {
	// LLM grammar suggestion covers "the the cat", overlapping the mechanics
	// "the the" finding at [0,7].
	client := &stubClient{response: `{"suggestions":[
		{"original":"the the cat","replacement":"the cat","explanation":"wordy","type":"grammar"}
	]}`}
	srv, docID, h := newTestServer(t, client)
	if _, err := h.DB.Exec(`UPDATE documents SET content_text = ? WHERE id = ?`, "the the cat sat", docID); err != nil {
		t.Fatalf("set content: %v", err)
	}

	// Persist the mechanics finding, then run the grammar pass.
	postMechanics(t, srv, docID, `[{"from":0,"to":7,"original":"the the","replacement":"the","explanation":"doubled word"}]`)
	rec := do(t, srv, http.MethodPost, "/docs/"+docID+"/check", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("check: code=%d body=%s", rec.Code, rec.Body)
	}
	var got []db.Suggestion
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode: %v", err)
	}
	for _, s := range got {
		if s.Type == db.SuggestionTypeGrammar {
			t.Fatalf("overlapping grammar card should have been dropped in favor of mechanics, got %+v", got)
		}
	}
	if len(got) != 1 || got[0].Type != db.SuggestionTypeMechanics {
		t.Fatalf("want sole mechanics finding, got %+v", got)
	}
}

// TestMechanicsActionedSuppression proves a dismissed mechanics fix is not
// re-proposed on the next submission of the same finding.
func TestMechanicsActionedSuppression(t *testing.T) {
	client := &stubClient{response: `{"suggestions":[]}`}
	srv, docID, _ := newTestServer(t, client)

	findings := `[{"from":6,"to":13,"original":"the the","replacement":"the","explanation":"doubled word"}]`
	got := postMechanics(t, srv, docID, findings)
	if len(got) != 1 {
		t.Fatalf("want 1 finding, got %+v", got)
	}
	// Dismiss it, then resubmit the same finding — it must stay suppressed.
	do(t, srv, http.MethodPost, "/suggestions/"+got[0].ID+"/dismiss", "")
	again := postMechanics(t, srv, docID, findings)
	if len(again) != 0 {
		t.Fatalf("dismissed mechanics fix should not reappear, got %+v", again)
	}
}