From f13a4533deb825ca98157558a8a376a288ebb109 Mon Sep 17 00:00:00 2001 From: Henry Date: Mon, 31 Aug 2026 21:50:17 +0000 Subject: [PATCH] perf(mcp): stop task-plan write guard from loading full revision history for one number (#3204) Co-authored-by: nova28 <17953305+nova28@users.noreply.github.com> Co-authored-by: Claude Sonnet 5 Co-authored-by: Carlos Florencio --- .../internal/mcp/handlers/task_plan_guard.go | 18 +++--- .../mcp/handlers/task_plan_guard_test.go | 41 +++++++++---- .../internal/task/service/plan_service.go | 17 +++++- .../task/service/plan_service_test.go | 58 +++++++++++++++++++ .../task/service/service_access_test.go | 6 ++ 5 files changed, 120 insertions(+), 20 deletions(-) diff --git a/apps/backend/internal/mcp/handlers/task_plan_guard.go b/apps/backend/internal/mcp/handlers/task_plan_guard.go index d565085672..3cab609281 100644 --- a/apps/backend/internal/mcp/handlers/task_plan_guard.go +++ b/apps/backend/internal/mcp/handlers/task_plan_guard.go @@ -107,13 +107,13 @@ type planWriteGuardResult struct { // A failure to fetch the current plan is non-fatal: the write proceeds // without a truncation warning, but it forces a new revision so a later // successful write cannot coalesce into an unknown prior revision. A failure -// to list the revision history is handled differently: -// truncation has already been detected from the plan content at that point, -// so the guard still forces a new revision and still warns — it renders the -// warning without a specific revision number (see planTruncationWarning) -// instead of silently dropping the warning, because clearing -// forceNewRevision here would let this destructive write coalesce into, and -// overwrite, the only surviving copy of the pre-truncation content. +// to fetch the latest revision is handled differently: truncation has +// already been detected from the plan content at that point, so the guard +// still forces a new revision and still warns — it renders the warning +// without a specific revision number (see planTruncationWarning) instead of +// silently dropping the warning, because clearing forceNewRevision here +// would let this destructive write coalesce into, and overwrite, the only +// surviving copy of the pre-truncation content. func (h *Handlers) evaluatePlanWriteGuard(ctx context.Context, taskID, newContent string) planWriteGuardResult { existing, err := h.planService.GetPlan(ctx, taskID) if err != nil { @@ -127,8 +127,8 @@ func (h *Handlers) evaluatePlanWriteGuard(ctx context.Context, taskID, newConten } priorRevisionNumber := 0 - if revisions, revErr := h.planService.ListRevisions(ctx, taskID); revErr == nil && len(revisions) > 0 { - priorRevisionNumber = revisions[0].RevisionNumber + if latest, revErr := h.planService.GetLatestRevision(ctx, taskID); revErr == nil && latest != nil { + priorRevisionNumber = latest.RevisionNumber } return planWriteGuardResult{ diff --git a/apps/backend/internal/mcp/handlers/task_plan_guard_test.go b/apps/backend/internal/mcp/handlers/task_plan_guard_test.go index 8045ab8e6b..3fb1ddd141 100644 --- a/apps/backend/internal/mcp/handlers/task_plan_guard_test.go +++ b/apps/backend/internal/mcp/handlers/task_plan_guard_test.go @@ -256,17 +256,29 @@ func TestMCPPlanTruncationGuard_CreateOverExistingPlanWarns(t *testing.T) { } } -// failingRevisionsRepo wraps the real sqlite repository but forces -// ListTaskPlanRevisions to always fail, simulating a transient lookup +// failingRevisionsRepo wraps the real sqlite repository but can force +// GetLatestTaskPlanRevision to fail on demand, simulating a transient lookup // failure (e.g. SQLITE_BUSY) that happens after GetPlan has already -// succeeded. +// succeeded. This is the call evaluatePlanWriteGuard actually makes to learn +// the prior revision number (via PlanService.GetLatestRevision). +// +// GetLatestTaskPlanRevision is also called by upsertPlan itself (to decide +// whether to coalesce) on every create/update, so failGetLatestRevision is a +// single-shot flag rather than an unconditional failure: tests arm it +// immediately before the call they want to fail, and it self-clears so the +// write's own internal lookup (which runs after the guard's) still succeeds. type failingRevisionsRepo struct { *sqlite.Repository - failGetPlan bool + failGetPlan bool + failGetLatestRevision bool } -func (r *failingRevisionsRepo) ListTaskPlanRevisions(context.Context, string, int) ([]*models.TaskPlanRevision, error) { - return nil, errors.New("simulated revision lookup failure") +func (r *failingRevisionsRepo) GetLatestTaskPlanRevision(ctx context.Context, taskID string) (*models.TaskPlanRevision, error) { + if r.failGetLatestRevision { + r.failGetLatestRevision = false + return nil, errors.New("simulated revision lookup failure") + } + return r.Repository.GetLatestTaskPlanRevision(ctx, taskID) } func (r *failingRevisionsRepo) GetTaskPlan(ctx context.Context, taskID string) (*models.TaskPlan, error) { @@ -278,8 +290,8 @@ func (r *failingRevisionsRepo) GetTaskPlan(ctx context.Context, taskID string) ( } // newMCPPlanTestHandlersWithFailingRevisionLookup builds Handlers like -// newMCPPlanTestHandlers, but backed by a plan service whose ListRevisions -// call always fails. +// newMCPPlanTestHandlers, but backed by a plan service whose +// GetLatestRevision call can be armed to fail once (single-shot). func newMCPPlanTestHandlersWithFailingRevisionLookup(t *testing.T) (*Handlers, *failingRevisionsRepo) { t.Helper() dbConn, err := db.OpenSQLite(filepath.Join(t.TempDir(), "test.db")) @@ -338,14 +350,14 @@ func newMCPPlanTestHandlersWithFailingRevisionLookup(t *testing.T) (*Handlers, * // TestMCPPlanTruncationGuard_RevisionLookupFailureOmitsRevisionNumber pins // Review round 2 Finding 3: revision numbering starts at 1 // (NextTaskPlanRevisionNumber), so "plan revision 0" can never be a real -// revision. When ListRevisions fails after GetPlan has already detected a +// revision. When GetLatestRevision fails after GetPlan has already detected a // truncating write, evaluatePlanWriteGuard must still force a new revision // and still warn — silently returning an empty guard here would let the // coalescing path overwrite the only surviving copy of the pre-truncation // content — but the rendered warning must not claim the content lives in // revision 0. func TestMCPPlanTruncationGuard_RevisionLookupFailureOmitsRevisionNumber(t *testing.T) { - h, _ := newMCPPlanTestHandlersWithFailingRevisionLookup(t) + h, repo := newMCPPlanTestHandlersWithFailingRevisionLookup(t) ctx := context.Background() large := strings.Repeat("x", 40000) @@ -362,6 +374,7 @@ func TestMCPPlanTruncationGuard_RevisionLookupFailureOmitsRevisionNumber(t *test // Check the guard directly before the write mutates the stored plan, so // "existing" below is still the large content. + repo.failGetLatestRevision = true guard := h.evaluatePlanWriteGuard(ctx, mcpPlanTaskID, small) if !guard.forceNewRevision { t.Error("forceNewRevision must stay true even when the revision lookup fails, " + @@ -373,7 +386,15 @@ func TestMCPPlanTruncationGuard_RevisionLookupFailureOmitsRevisionNumber(t *test if strings.Contains(guard.warning, "revision 0") { t.Errorf("warning names a nonexistent revision 0 (revision numbering starts at 1): %q", guard.warning) } + if guard.priorRevision != 0 { + t.Errorf("priorRevision must be 0 when the revision lookup fails, got %d", guard.priorRevision) + } + // Arm the failure again for the guard's own lookup inside the write path. + // It self-clears after firing once, so the write's own internal + // GetLatestTaskPlanRevision call (upsertPlan's coalesce decision) still + // succeeds and the update itself is not blocked. + repo.failGetLatestRevision = true out, err := h.handleUpdateTaskPlan(ctx, mcpPlanMsg(t, ws.ActionMCPUpdateTaskPlan, mustMarshalPlanPayload(t, map[string]any{ "task_id": mcpPlanTaskID, diff --git a/apps/backend/internal/task/service/plan_service.go b/apps/backend/internal/task/service/plan_service.go index 70dd6391f3..e65cee6653 100644 --- a/apps/backend/internal/task/service/plan_service.go +++ b/apps/backend/internal/task/service/plan_service.go @@ -410,7 +410,8 @@ func (s *PlanService) DeletePlan(ctx context.Context, taskID string) error { return nil } -// ListRevisions returns plan revisions newest-first without content (metadata only). +// ListRevisions returns every plan revision for a task, newest-first, each +// including its full content. func (s *PlanService) ListRevisions(ctx context.Context, taskID string) ([]*models.TaskPlanRevision, error) { if taskID == "" { return nil, ErrTaskIDRequired @@ -421,6 +422,20 @@ func (s *PlanService) ListRevisions(ctx context.Context, taskID string) ([]*mode return s.repo.ListTaskPlanRevisions(ctx, taskID, 0) } +// GetLatestRevision returns the most recent revision for a task, or nil if +// none exist. Unlike ListRevisions, it fetches only the latest revision row +// instead of every revision — callers that only need the latest revision +// number (e.g. the plan write guard) should use this instead. +func (s *PlanService) GetLatestRevision(ctx context.Context, taskID string) (*models.TaskPlanRevision, error) { + if taskID == "" { + return nil, ErrTaskIDRequired + } + if err := s.authorize(ctx, taskID); err != nil { + return nil, err + } + return s.repo.GetLatestTaskPlanRevision(ctx, taskID) +} + // GetRevision returns a single revision with content (for diff/preview). func (s *PlanService) GetRevision(ctx context.Context, id string) (*models.TaskPlanRevision, error) { rev, err := s.repo.GetTaskPlanRevision(ctx, id) diff --git a/apps/backend/internal/task/service/plan_service_test.go b/apps/backend/internal/task/service/plan_service_test.go index b52541c3ef..9e51bba2e2 100644 --- a/apps/backend/internal/task/service/plan_service_test.go +++ b/apps/backend/internal/task/service/plan_service_test.go @@ -439,6 +439,64 @@ func TestPlanService_AppendsWhenWindowExpired(t *testing.T) { } } +// TestPlanService_GetLatestRevision pins the cheap latest-revision accessor: +// it must report the same revision number ListRevisions()[0] would, without +// requiring callers to load every revision's content just to read one int. +func TestPlanService_GetLatestRevision(t *testing.T) { + svc, _, repo := createTestPlanService(t) + ctx := context.Background() + seedTask(t, ctx, repo, "task-latest") + svc.coalesceWindow = 0 // disable coalescing so each write is a new revision + + // No plan yet: nil, nil. + latest, err := svc.GetLatestRevision(ctx, "task-latest") + if err != nil { + t.Fatalf("GetLatestRevision (no plan): %v", err) + } + if latest != nil { + t.Errorf("expected nil latest revision before any plan exists, got %+v", latest) + } + + _, _ = svc.CreatePlan(ctx, CreatePlanRequest{ + TaskID: "task-latest", Content: "v1", + AuthorKind: "agent", AuthorName: "Claude", + }) + _, _ = svc.CreatePlan(ctx, CreatePlanRequest{ + TaskID: "task-latest", Content: "v2", + AuthorKind: "agent", AuthorName: "Claude", + }) + + latest, err = svc.GetLatestRevision(ctx, "task-latest") + if err != nil { + t.Fatalf("GetLatestRevision: %v", err) + } + if latest == nil { + t.Fatal("expected a latest revision after two writes, got nil") + } + if latest.RevisionNumber != 2 || latest.Content != "v2" { + t.Errorf("latest = %+v, want revision 2 with content v2", latest) + } + + list, _ := svc.ListRevisions(ctx, "task-latest") + if len(list) == 0 { + t.Fatal("GetLatestRevision: ListRevisions returned no revisions after two writes") + } + if list[0].RevisionNumber != latest.RevisionNumber { + t.Errorf("GetLatestRevision disagrees with ListRevisions()[0]: got %d, want %d", + latest.RevisionNumber, list[0].RevisionNumber) + } +} + +func TestPlanService_GetLatestRevisionRequiresTaskID(t *testing.T) { + svc, _, _ := createTestPlanService(t) + ctx := context.Background() + + _, err := svc.GetLatestRevision(ctx, "") + if err != ErrTaskIDRequired { + t.Errorf("expected ErrTaskIDRequired, got %v", err) + } +} + func TestPlanService_ForceNewRevisionPreventsCoalesce(t *testing.T) { svc, _, repo := createTestPlanService(t) ctx := context.Background() diff --git a/apps/backend/internal/task/service/service_access_test.go b/apps/backend/internal/task/service/service_access_test.go index b671611b4c..c173d305a0 100644 --- a/apps/backend/internal/task/service/service_access_test.go +++ b/apps/backend/internal/task/service/service_access_test.go @@ -345,6 +345,12 @@ func TestPlanAndWalkthroughScoping(t *testing.T) { if _, err := plan.GetPlan(ctxAs("user-b"), "task-b"); errors.Is(err, repoerrors.ErrTaskNotFound) { t.Fatalf("owner plan get must pass the auth gate: %v", err) } + if _, err := plan.GetLatestRevision(ctxAs("user-a"), "task-b"); !errors.Is(err, repoerrors.ErrTaskNotFound) { + t.Fatalf("foreign latest revision get: %v", err) + } + if _, err := plan.GetLatestRevision(ctxAs("user-b"), "task-b"); errors.Is(err, repoerrors.ErrTaskNotFound) { + t.Fatalf("owner latest revision get must pass the auth gate: %v", err) + } wt := NewWalkthroughService(repo, nil, log) wt.SetTaskAuthorizer(svc.AuthorizeTaskAccess)