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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user