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 <noreply@anthropic.com>
Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>
This commit is contained in:
Henry
2026-08-31 21:50:17 +00:00
committed by GitHub
parent b292f9241e
commit f13a4533de
5 changed files with 120 additions and 20 deletions
@@ -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{
@@ -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,
@@ -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)
@@ -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()
@@ -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)