fix: settle approval and worktree races

pkg/aichat/approval_execution.go

@@ -11,30 +11,39 @@ import (
 )
 
 const (
-	suspendedSeedTimeout  = 15 * time.Second
+	// suspendedSeedWait bounds how long an approval resolution waits for the
+	// suspending turn's assistant message to land in the thread store.
+	suspendedSeedWait     = 5 * time.Second
 	suspendedSeedInterval = 25 * time.Millisecond
 )
 
-// awaitSuspendedSeed returns the assistant message the suspended turn ended on.
-// The durable suspension (prompt run -> waiting) is committed while the same
-// stream's persistence goroutine is still writing that assistant message, so an
-// approval resolved the instant the run becomes resumable can observe the run
-// before its transcript. Wait for that in-flight write instead of rejecting a
-// legitimate approval, and fail loudly when it never lands.
+// awaitSuspendedSeed returns the thread's trailing assistant message for the
+// suspended turn.
+//
+// The durable suspension (prompt run -> waiting) is committed from the event
+// pipeline while the same stream is still persisting the assistant message it
+// suspended on, so an approval resolved the instant the run becomes resumable
+// can observe the run before its transcript. The run's durable approval state
+// guarantees that message is committed or imminent — wait the in-flight write
+// out instead of failing a resolution that has already consumed the approval,
+// and fail loudly when it never lands.
 func awaitSuspendedSeed(ctx context.Context, store ThreadStore, threadID, turnID string) (*UIMessage, error) {
-	deadline := time.Now().Add(suspendedSeedTimeout)
+	deadline := time.Now().Add(suspendedSeedWait)
 	for {
 		thread, err := store.Get(ctx, threadID)
 		if err != nil {
 			return nil, err
 		}
-		if len(thread.Messages) > 0 {
-			seed := thread.Messages[len(thread.Messages)-1]
+		if count := len(thread.Messages); count > 0 {
+			seed := thread.Messages[count-1]
 			if strings.EqualFold(seed.Role, string(api.RoleAssistant)) && seed.TurnID == turnID {
 				return &seed, nil
 			}
 		}
 		if time.Now().After(deadline) {
+			if len(thread.Messages) == 0 {
+				return nil, fmt.Errorf("captain chat session %s has no suspended assistant message", threadID)
+			}
 			return nil, fmt.Errorf("captain chat session %s does not end with the suspended turn %s", threadID, turnID)
 		}
 		select {
@@ -155,44 +164,11 @@ func enforceApprovalRuntimeProfile(spec api.Spec, resolved api.ComposedSpec) err
 	return nil
 }
 
-// suspendedSeedWait bounds how long an approval resolution waits for the
-// suspending turn's assistant message to land in the thread store.
-const suspendedSeedWait = 5 * time.Second
-
-// awaitSuspendedSeed returns the thread's trailing assistant message for the
-// suspended turn. The prompt run reaches its waiting state from the event
-// pipeline before the suspending stream persists that message on its final
-// unwind, so an approval resolved from a session poll can arrive while the
-// write is still in flight. The run's durable approval state guarantees the
-// message is committed or imminent — wait it out instead of failing a
-// resolution that has already consumed the approval.
+// awaitSuspendedSeed resolves this service's thread store and waits there.
 func (s *Service) awaitSuspendedSeed(ctx context.Context, threadID, turnID string) (*UIMessage, error) {
 	store, err := s.threads(ctx)
 	if err != nil {
 		return nil, err
 	}
-	deadline := time.Now().Add(suspendedSeedWait)
-	for {
-		thread, err := store.Get(ctx, threadID)
-		if err != nil {
-			return nil, err
-		}
-		if count := len(thread.Messages); count > 0 {
-			seed := thread.Messages[count-1]
-			if strings.EqualFold(seed.Role, string(api.RoleAssistant)) && seed.TurnID == turnID {
-				return &seed, nil
-			}
-		}
-		if time.Now().After(deadline) {
-			if len(thread.Messages) == 0 {
-				return nil, fmt.Errorf("captain chat session %s has no suspended assistant message", threadID)
-			}
-			return nil, fmt.Errorf("captain chat session %s does not end with the suspended turn %s", threadID, turnID)
-		}
-		select {
-		case <-ctx.Done():
-			return nil, ctx.Err()
-		case <-time.After(50 * time.Millisecond):
-		}
-	}
+	return awaitSuspendedSeed(ctx, store, threadID, turnID)
 }

pkg/aichat/approval_settle_integration_test.go

@@ -0,0 +1,149 @@
+package aichat_test
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"github.com/flanksource/captain/pkg/aichat"
+	"github.com/flanksource/captain/pkg/api"
+	"github.com/flanksource/captain/pkg/database"
+	"github.com/flanksource/commons-db/dbtest"
+	"github.com/google/uuid"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+// A provider approval is answerable — it is on the session, and its ID has gone
+// out on the event stream — from the moment the permission frame is observed.
+// The run it blocks only reaches `waiting` once the stream has finished the turn
+// and encoded its checkpoint, several statements later. These specs pin what
+// happens to an answer that arrives inside that window, which is where a person
+// clicking Approve promptly, and the mocked lifecycle suite, both landed.
+var _ = Describe("Approvals answered while the suspension is still landing", func() {
+	It("waits for the run to park rather than refusing the answer", func(ctx SpecContext) {
+		fixture := newApprovalFixture(ctx, "captain_aichat_approval_settles")
+
+		// Nothing has parked the run yet; under the bare store guard this is the
+		// exact moment that produced "cannot be resolved before its prompt run is
+		// waiting" and a 409.
+		suspended := make(chan error, 1)
+		go func() {
+			time.Sleep(200 * time.Millisecond)
+			suspended <- suspendOnAccountsApproval(ctx, fixture.execution)
+		}()
+
+		continuation, err := fixture.authority.ResolveToolApproval(ctx, aichat.ToolApprovalResolution{
+			ThreadID: fixture.thread.ID, ApprovalID: fixture.approvalID, Approved: false, Reason: "not now",
+		})
+		Expect(err).NotTo(HaveOccurred(), "an answer that raced the suspension is still a valid answer")
+		Expect(continuation).NotTo(BeNil(), "the resolution has to hand back the continuation that resumes the run")
+		DeferCleanup(continuation.Execution.Close)
+		Expect(<-suspended).To(Succeed())
+
+		resolved, err := fixture.store.GetSession(ctx, fixture.thread.ID)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(resolved.Requests).To(HaveLen(1))
+		Expect(resolved.Requests[0].State).To(Equal(string(database.TurnRequestStateDenied)))
+		Expect(resolved.Requests[0].Reason).To(Equal("not now"))
+	})
+
+	It("refuses an answer once the run it blocks has already ended", func(ctx SpecContext) {
+		fixture := newApprovalFixture(ctx, "captain_aichat_approval_run_ended")
+
+		runID, err := uuid.Parse(fixture.execution.PromptRunID())
+		Expect(err).NotTo(HaveOccurred())
+		run, err := fixture.db.GetPromptRun(ctx, runID)
+		Expect(err).NotTo(HaveOccurred())
+		cancelled := database.PromptRunStateCancelled
+		_, err = fixture.db.UpdatePromptRun(ctx, database.UpdatePromptRunInput{
+			ID: run.ID, ExpectedVersion: run.Version, State: &cancelled,
+		})
+		Expect(err).NotTo(HaveOccurred())
+
+		// No suspension is coming, so this must fail on the run's state rather
+		// than burn the whole settle budget waiting for one.
+		started := time.Now()
+		_, err = fixture.authority.ResolveToolApproval(ctx, aichat.ToolApprovalResolution{
+			ThreadID: fixture.thread.ID, ApprovalID: fixture.approvalID, Approved: true,
+		})
+		Expect(err).To(MatchError(database.ErrTurnRequestConflict))
+		Expect(err).To(MatchError(ContainSubstring("already ended (cancelled)")))
+		Expect(time.Since(started)).To(BeNumerically("<", time.Second),
+			"a run that ended is a decided answer, not something to wait out")
+	})
+})
+
+type approvalFixture struct {
+	db         *database.DB
+	store      *aichat.DatabaseThreadStore
+	authority  *aichat.DatabaseExecutionAuthority
+	thread     *aichat.Thread
+	execution  aichat.Execution
+	approvalID string
+}
+
+// newApprovalFixture drives a chat turn up to the point where the provider has
+// asked for permission and the durable approval exists, but the run has not yet
+// been parked.
+func newApprovalFixture(ctx context.Context, name string) approvalFixture {
+	GinkgoHelper()
+	testDB := dbtest.ForGinkgo(dbtest.Options{Name: name})
+	db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations())
+	Expect(err).NotTo(HaveOccurred())
+	DeferCleanup(db.Close)
+	store, err := aichat.NewDatabaseThreadStore(db)
+	Expect(err).NotTo(HaveOccurred())
+	thread, err := store.Create(ctx, "Accounts")
+	Expect(err).NotTo(HaveOccurred())
+	authority, err := aichat.NewDatabaseExecutionAuthority(db)
+	Expect(err).NotTo(HaveOccurred())
+	execution, err := authority.Begin(ctx, aichat.ExecutionRequest{
+		ThreadID: thread.ID, RequestID: "user-message-1", Title: thread.Title,
+		Spec: api.Spec{Model: withCaps(api.Model{Name: "gemini", Mode: api.ModeAPI})},
+	})
+	Expect(err).NotTo(HaveOccurred())
+	Expect(store.AppendMessage(ctx, thread.ID, aichat.UIMessage{
+		ID: "user-message-1", TurnID: execution.TurnID(), Role: "user",
+		Parts: []aichat.UIPart{{Type: "text", Text: "Edit the account"}},
+	})).To(Succeed())
+	permission, err := execution.Observe(ctx, api.Event{
+		Kind: api.EventPermission, ToolCallID: "call-account-1", Tool: "accounts_edit",
+		Input: map[string]any{"id": "acc-1"},
+	})
+	Expect(err).NotTo(HaveOccurred())
+	Expect(store.AppendMessage(ctx, thread.ID, aichat.UIMessage{
+		ID: execution.TurnID() + "-assistant", TurnID: execution.TurnID(), Role: "assistant",
+		Parts: []aichat.UIPart{{
+			Type: "dynamic-tool", ToolName: "accounts_edit", ToolCallID: "call-account-1",
+			State: "approval-requested", Input: json.RawMessage(`{"id":"acc-1"}`),
+			Approval: &aichat.Approval{ID: permission.ApprovalID},
+		}},
+	})).To(Succeed())
+	return approvalFixture{
+		db: db, store: store, authority: authority, thread: thread,
+		execution: execution, approvalID: permission.ApprovalID,
+	}
+}
+
+// suspendOnAccountsApproval completes the turn the way a provider that needs an
+// approval does: a terminal result carrying the approval state and the private
+// checkpoint the resume replays from. This is what parks the run in `waiting`.
+func suspendOnAccountsApproval(ctx context.Context, execution aichat.Execution) error {
+	_, err := execution.Observe(ctx, api.Event{
+		Kind: api.EventResult, Success: true,
+		ToolApproval: &api.ToolApprovalState{
+			Messages: []api.Message{{Role: api.RoleAssistant, Parts: []api.Part{{
+				Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{
+					ToolCallID: "call-account-1", Name: "accounts_edit", Input: json.RawMessage(`{"id":"acc-1"}`),
+				},
+			}}}},
+			Calls: []api.ToolApprovalCall{{Request: api.ToolApprovalRequest{
+				ToolCallID: "call-account-1", Tool: "accounts_edit", Input: json.RawMessage(`{"id":"acc-1"}`),
+			}}},
+			ProviderCheckpoint: &api.ProviderCheckpoint{Codec: "test-provider", Version: 1, Payload: []byte("checkpoint")},
+		},
+	})
+	return err
+}

pkg/aichat/database_threads_integration_test.go

@@ -218,8 +218,6 @@ var _ = Describe("Database chat sessions", func() {
 		resolution := aichat.ToolApprovalResolution{
 			ThreadID: thread.ID, ApprovalID: permission.ApprovalID, Approved: false, Reason: "not now",
 		}
-		_, err = authority.ResolveToolApproval(ctx, resolution)
-		Expect(err).To(MatchError(ContainSubstring("cannot be resolved before its prompt run is waiting")))
 		assistant := aichat.UIMessage{
 			ID: execution.TurnID() + "-assistant", TurnID: execution.TurnID(), Role: "assistant",
 			Parts: []aichat.UIPart{{
@@ -229,21 +227,7 @@ var _ = Describe("Database chat sessions", func() {
 			}},
 		}
 		Expect(store.AppendMessage(ctx, thread.ID, assistant)).To(Succeed())
-		_, err = execution.Observe(ctx, api.Event{
-			Kind: api.EventResult, Success: true,
-			ToolApproval: &api.ToolApprovalState{
-				Messages: []api.Message{{Role: api.RoleAssistant, Parts: []api.Part{{
-					Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{
-						ToolCallID: "call-account-1", Name: "accounts_edit", Input: json.RawMessage(`{"id":"acc-1"}`),
-					},
-				}}}},
-				Calls: []api.ToolApprovalCall{{Request: api.ToolApprovalRequest{
-					ToolCallID: "call-account-1", Tool: "accounts_edit", Input: json.RawMessage(`{"id":"acc-1"}`),
-				}}},
-				ProviderCheckpoint: &api.ProviderCheckpoint{Codec: "test-provider", Version: 1, Payload: []byte("checkpoint")},
-			},
-		})
-		Expect(err).NotTo(HaveOccurred())
+		Expect(suspendOnAccountsApproval(ctx, execution)).To(Succeed())
 
 		aggregate, err := store.GetSession(ctx, thread.ID)
 		Expect(err).NotTo(HaveOccurred())

pkg/aichat/execution_database_authority.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"maps"
 	"strings"
+	"time"
 
 	"github.com/flanksource/captain/pkg/api"
 	"github.com/flanksource/captain/pkg/database"
@@ -132,10 +133,82 @@ func (a *DatabaseExecutionAuthority) Begin(
 	return execution, nil
 }
 
+const (
+	// suspendedRunWait bounds how long a resolution waits for the run it answers
+	// to finish parking, and suspendedRunInterval is how often that is re-read.
+	// They match the seed wait's budget: the two wait out the two halves of the
+	// same in-flight suspension.
+	suspendedRunWait     = 5 * time.Second
+	suspendedRunInterval = 25 * time.Millisecond
+)
+
+// awaitSuspendedRun waits for a provider approval's prompt run to reach
+// `waiting`, the one state ResolveToolApprovalRequest accepts an answer in.
+//
+// A provider approval is recorded — and so becomes visible on the session and
+// goes out on the event stream carrying its approval ID — while the stream that
+// raised it is still finishing the turn and encoding its checkpoint. The run
+// only reaches `waiting` several statements later. So anything that answers the
+// question the moment it is asked raced the suspension and got a 409 telling it
+// to retry something that was never wrong: a person clicking Approve promptly,
+// or a poller in a test.
+//
+// The guard being waited for is not removable. A resolution applied before the
+// run parks yields no continuation (see resolveToolApproval), and the suspension
+// then parks the run on an already-answered approval that nothing ever resumes.
+// So wait the parking out — the same treatment awaitSuspendedSeed gives the
+// other half of this window — and fail loudly when it never happens.
+//
+// This runs outside the resolving transaction deliberately: a snapshot taken
+// inside one would never observe the suspending connection's commit.
+func (a *DatabaseExecutionAuthority) awaitSuspendedRun(ctx context.Context, approvalID string) error {
+	requestID, err := uuid.Parse(approvalID)
+	if err != nil {
+		return nil // resolveToolApproval reports a malformed ID, with its own message
+	}
+	deadline := time.Now().Add(suspendedRunWait)
+	for {
+		request, err := a.db.GetTurnRequest(ctx, requestID)
+		if err != nil {
+			return nil // the resolve path owns not-found and read failures alike
+		}
+		// A caller-tool approval carries its own authority and is answerable
+		// whatever its run is doing. Anything already decided, or with no run to
+		// resume, is likewise the store's answer to give, not this wait's.
+		if request.CredentialID != nil || request.PromptRunID == nil ||
+			request.State != database.TurnRequestStatePending {
+			return nil
+		}
+		run, err := a.db.GetPromptRun(ctx, *request.PromptRunID)
+		if err != nil {
+			return err
+		}
+		switch run.State {
+		case database.PromptRunStateWaiting:
+			return nil
+		case database.PromptRunStateSucceeded, database.PromptRunStateFailed, database.PromptRunStateCancelled:
+			return fmt.Errorf("%w: approval %s cannot be resolved, its prompt run %s already ended (%s)",
+				database.ErrTurnRequestConflict, request.ID, run.ID, run.State)
+		}
+		if time.Now().After(deadline) {
+			return fmt.Errorf("%w: approval %s is still pending after %s with its prompt run %s in state %q rather than waiting",
+				database.ErrTurnRequestConflict, request.ID, suspendedRunWait, run.ID, run.State)
+		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(suspendedRunInterval):
+		}
+	}
+}
+
 func (a *DatabaseExecutionAuthority) ResolveToolApproval(
 	ctx context.Context,
 	resolution ToolApprovalResolution,
 ) (*ApprovalContinuation, error) {
+	if err := a.awaitSuspendedRun(ctx, resolution.ApprovalID); err != nil {
+		return nil, err
+	}
 	var continuation *ApprovalContinuation
 	err := a.db.Transaction(ctx, func(tx *database.DB) error {
 		var resolveErr error

pkg/aichat/execution_database_integration_test.go

@@ -385,8 +385,13 @@ var _ = Describe("Database execution authority", func() {
 		continueResolution := make(chan struct{})
 		var intercepted atomic.Bool
 		const callback = "test:pause_approval_after_prompt_run_read"
+		// Pause the run read the resolution *resumes* from — the one inside its
+		// transaction, whose version the update then asserts on. Resolution also
+		// reads the run outside any transaction first, to wait out a suspension
+		// still landing; pausing there would stall a read this race is not about.
 		Expect(db.Gorm().Callback().Query().After("gorm:query").Register(callback, func(tx *gorm.DB) {
-			if tx.Statement.Table == "captain_prompt_runs" && intercepted.CompareAndSwap(false, true) {
+			_, inTransaction := tx.Statement.ConnPool.(gorm.TxCommitter)
+			if inTransaction && tx.Statement.Table == "captain_prompt_runs" && intercepted.CompareAndSwap(false, true) {
 				close(versionRead)
 				<-continueResolution
 			}