fix(orchestrator): plug zombie turn leak pinning sessions to RUNNING (#712)

* fix(orchestrator): plug zombie turn leak pinning sessions to RUNNING Two independent code paths lazily created turns without sharing state. service.CreateMessage (DB-based) created a turn for inbound user messages; PromptTask then unconditionally created another via startTurnForSession (in-memory + DB). The first turn was never tracked by the orchestrator, so its completed_at never got set and the session state machine kept the session pinned to RUNNING forever — the "agent is processing" spinner that wouldn't clear. - startTurnForSession is now idempotent: adopts an existing open turn from the in-memory map or the DB before creating a new one. - completeTurnForSession is now DB-authoritative: the in-memory map is a hint, the DB is truth. Loops to mop up multiple pre-existing zombies with a sanity bound. - CancelAgent captures the turn ID before agentManager.CancelAgent so the cancel message can't race the agent's complete event into a phantom turn. - service.GetActiveTurn returns (nil, nil) for "no active turn" instead of leaking sql.ErrNoRows. Adds five regression tests covering adoption, idempotency, untracked close, multi-zombie sweep, and double-complete idempotency. * fix(orchestrator): address PR review feedback on turn lifecycle - Log DB read errors in startTurnForSession adoption path so transient failures aren't silently swallowed (Claude, CodeRabbit, Greptile). - completeTurnForSession probes once after the iteration cap and only warns when a turn is *still* observed — closing exactly maxIterations is no longer a false alarm (CodeRabbit). - Clarify the CompleteTurn-failure log: returning early defers cleanup to the next sweep rather than silently dropping zombies (Greptile). - Route getOrStartTurn through GetActiveTurn so ErrNoRows normalization is the single contract (Claude). Adds TestCompleteTurnRespectsIterationCap to lock the cap behavior. * fix(orchestrator): align getOrStartTurn observability with adoption path Log DB read errors in getOrStartTurn before falling through to StartTurn — same observability contract as the orchestrator's startTurnForSession (Claude review on the previous fix-up). Also add a comment in turn_lifecycle_test.go explaining the hardcoded "task1" TaskID in repoTurnService matches seedSession, so future test authors reusing the helper know the coupling. --------- Co-authored-by: Kandev Agent <agent@kandev.dev>
2026-04-27 17:41:22 +01:00
parent b805bf796a
commit 8bf97911d5
4 changed files with 326 additions and 21 deletions
@@ -459,12 +459,41 @@ func (s *Service) initWorkflowEngine() {
 	s.workflowEngine = engine.New(store, callbacks)
 }

-// startTurnForSession starts a new turn for the session and stores it.
+// startTurnForSession ensures the session has an active turn and returns its ID.
+//
+// Idempotent: if an open turn already exists (either tracked in activeTurns or
+// only present in the DB — the latter happens when service.CreateMessage lazily
+// started a turn for an inbound user message before the orchestrator's prompt
+// cycle began, or when the backend restarted with open turns in the DB), it is
+// adopted rather than duplicated. A new turn is created only when none exists.
+//
+// This avoids the classic dual-creation leak: user message → service.CreateMessage
+// starts turn X (DB only) → PromptTask → startTurnForSession → would create turn Y
+// (DB + activeTurns), leaving X open forever because nothing tracks it.
 func (s *Service) startTurnForSession(ctx context.Context, sessionID string) string {
 	if s.turnService == nil {
 		return ""
 	}

+	if turnIDVal, ok := s.activeTurns.Load(sessionID); ok {
+		if turnID, ok := turnIDVal.(string); ok && turnID != "" {
+			return turnID
+		}
+	}
+
+	if turn, err := s.turnService.GetActiveTurn(ctx, sessionID); turn != nil {
+		s.activeTurns.Store(sessionID, turn.ID)
+		return turn.ID
+	} else if err != nil {
+		// A real DB read failure here would otherwise be silently dropped, and
+		// we'd fall through to StartTurn — potentially writing a duplicate next
+		// to an existing open turn we couldn't see. Log it; the next sweep via
+		// completeTurnForSession will mop up any duplicate.
+		s.logger.Warn("failed to look up active turn before starting a new one",
+			zap.String("session_id", sessionID),
+			zap.Error(err))
+	}
+
 	turn, err := s.turnService.StartTurn(ctx, sessionID)
 	if err != nil {
 		s.logger.Warn("failed to start turn",
@@ -477,27 +506,51 @@ func (s *Service) startTurnForSession(ctx context.Context, sessionID string) str
 	return turn.ID
 }

-// completeTurnForSession completes the active turn for the session.
+// completeTurnForSession closes any open turn for the session.
+//
+// The DB is the source of truth — activeTurns is just an in-memory cache that
+// can drift (see startTurnForSession) or be wiped by a backend restart. We
+// query the DB for any open turn and close it. Loops to mop up multiple
+// zombies (e.g. left over from before this fix) with a small sanity bound.
 func (s *Service) completeTurnForSession(ctx context.Context, sessionID string) {
 	if s.turnService == nil {
 		return
 	}

-	turnIDVal, ok := s.activeTurns.LoadAndDelete(sessionID)
-	if !ok {
-		return
-	}
+	s.activeTurns.Delete(sessionID)

-	turnID, ok := turnIDVal.(string)
-	if !ok || turnID == "" {
-		return
+	const maxIterations = 16
+	closed := 0
+	for closed < maxIterations {
+		turn, err := s.turnService.GetActiveTurn(ctx, sessionID)
+		if err != nil {
+			s.logger.Warn("failed to look up active turn",
+				zap.String("session_id", sessionID),
+				zap.Error(err))
+			return
+		}
+		if turn == nil {
+			return
+		}
+		if err := s.turnService.CompleteTurn(ctx, turn.ID); err != nil {
+			// GetActiveTurn returns the latest open turn — retrying here
+			// would just hit the same row and loop. Bail; the next
+			// completeTurnForSession call will pick it up.
+			s.logger.Warn("failed to complete turn; will retry on next sweep",
+				zap.String("session_id", sessionID),
+				zap.String("turn_id", turn.ID),
+				zap.Error(err))
+			return
+		}
+		closed++
 	}
-
-	if err := s.turnService.CompleteTurn(ctx, turnID); err != nil {
-		s.logger.Warn("failed to complete turn",
+	// Only warn if turns are *still* accumulating after the cap. Closing
+	// exactly maxIterations turns and then finding the session clean is not a
+	// runaway.
+	if turn, err := s.turnService.GetActiveTurn(ctx, sessionID); err == nil && turn != nil {
+		s.logger.Warn("completeTurnForSession iteration cap hit; possible turn close loop",
 			zap.String("session_id", sessionID),
-			zap.String("turn_id", turnID),
-			zap.Error(err))
+			zap.Int("max_iterations", maxIterations))
 	}
 }

@@ -1803,6 +1803,13 @@ func (s *Service) CancelAgent(ctx context.Context, sessionID string) error {
 			zap.Error(err))
 	}

+	// Capture the active turn before cancelling so the cancel message attaches
+	// to the turn the user was actually cancelling. If we waited until after
+	// agentManager.CancelAgent, the agent's complete event could have already
+	// closed the turn, and getActiveTurnID would lazily create a phantom turn
+	// just to host the cancel message.
+	cancelTurnID := s.getActiveTurnID(sessionID)
+
 	if err := s.agentManager.CancelAgent(ctx, sessionID); err != nil {
 		switch {
 		case errors.Is(err, lifecycle.ErrNoExecutionForSession):
@@ -1842,7 +1849,7 @@ func (s *Service) CancelAgent(ctx context.Context, sessionID string) error {
 			"Turn cancelled by user",
 			sessionID,
 			string(v1.MessageTypeStatus),
-			s.getActiveTurnID(sessionID),
+			cancelTurnID,
 			metadata,
 			false,
 		); err != nil {
@@ -1852,7 +1859,8 @@ func (s *Service) CancelAgent(ctx context.Context, sessionID string) error {
 		}
 	}

-	// Complete the turn since the agent was cancelled
+	// Complete the turn since the agent was cancelled. Idempotent w.r.t. a
+	// concurrent agent.complete event having already closed the turn.
 	s.completeTurnForSession(ctx, sessionID)

 	s.logger.Debug("agent turn cancelled", zap.String("session_id", sessionID))
@@ -0,0 +1,232 @@
+package orchestrator
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+
+	"github.com/kandev/kandev/internal/task/models"
+	sqliterepo "github.com/kandev/kandev/internal/task/repository/sqlite"
+)
+
+// repoTurnService is a minimal TurnService used by lifecycle tests. It mirrors
+// the production task service's behavior for the three methods the orchestrator
+// uses, while sharing the same sqlite repo as the rest of the test setup so
+// that DB-backed assertions stay coherent across components.
+type repoTurnService struct {
+	repo *sqliterepo.Repository
+}
+
+func (a *repoTurnService) StartTurn(ctx context.Context, sessionID string) (*models.Turn, error) {
+	now := time.Now().UTC()
+	turn := &models.Turn{
+		ID:            uuid.New().String(),
+		TaskSessionID: sessionID,
+		TaskID:        "task1", // matches the taskID seedSession uses
+		StartedAt:     now,
+		CreatedAt:     now,
+		UpdatedAt:     now,
+	}
+	if err := a.repo.CreateTurn(ctx, turn); err != nil {
+		return nil, err
+	}
+	return turn, nil
+}
+
+func (a *repoTurnService) CompleteTurn(ctx context.Context, turnID string) error {
+	return a.repo.CompleteTurn(ctx, turnID)
+}
+
+func (a *repoTurnService) GetActiveTurn(ctx context.Context, sessionID string) (*models.Turn, error) {
+	turn, err := a.repo.GetActiveTurnBySessionID(ctx, sessionID)
+	if errors.Is(err, sql.ErrNoRows) {
+		return nil, nil
+	}
+	return turn, err
+}
+
+func openTurnCount(t *testing.T, repo *sqliterepo.Repository, sessionID string) int {
+	t.Helper()
+	turns, err := repo.ListTurnsBySession(context.Background(), sessionID)
+	if err != nil {
+		t.Fatalf("ListTurnsBySession: %v", err)
+	}
+	open := 0
+	for _, turn := range turns {
+		if turn.CompletedAt == nil {
+			open++
+		}
+	}
+	return open
+}
+
+func newTurnLifecycleTestService(t *testing.T) (*Service, *sqliterepo.Repository) {
+	t.Helper()
+	repo := setupTestRepo(t)
+	seedSession(t, repo, "task1", "session1", "step1")
+	svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
+	svc.turnService = &repoTurnService{repo: repo}
+	return svc, repo
+}
+
+// TestStartTurnAdoptsExistingDBTurn covers the dual-creation leak that left
+// zombie turns whenever service.CreateMessage lazily started a turn for an
+// inbound user message and the orchestrator's PromptTask then started another.
+// Now startTurnForSession adopts the open DB turn instead of creating a second.
+func TestStartTurnAdoptsExistingDBTurn(t *testing.T) {
+	svc, repo := newTurnLifecycleTestService(t)
+	ctx := context.Background()
+
+	// Simulate service.CreateMessage having created a turn behind the
+	// orchestrator's back (DB row only, not in activeTurns).
+	preexisting, err := svc.turnService.StartTurn(ctx, "session1")
+	if err != nil {
+		t.Fatalf("seed turn: %v", err)
+	}
+
+	// PromptTask path calls startTurnForSession after setSessionRunning. With
+	// the fix it must adopt the preexisting DB turn rather than create another.
+	adopted := svc.startTurnForSession(ctx, "session1")
+	if adopted != preexisting.ID {
+		t.Fatalf("expected adoption of existing turn %q, got %q", preexisting.ID, adopted)
+	}
+
+	turns, err := repo.ListTurnsBySession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("ListTurnsBySession: %v", err)
+	}
+	if len(turns) != 1 {
+		t.Fatalf("expected 1 turn, got %d (zombies: %d)", len(turns), openTurnCount(t, repo, "session1"))
+	}
+}
+
+// TestStartTurnIsIdempotentInMemory verifies that repeated calls do not stack
+// turns when activeTurns already tracks one.
+func TestStartTurnIsIdempotentInMemory(t *testing.T) {
+	svc, repo := newTurnLifecycleTestService(t)
+	ctx := context.Background()
+
+	first := svc.startTurnForSession(ctx, "session1")
+	second := svc.startTurnForSession(ctx, "session1")
+	if first == "" {
+		t.Fatal("expected a turn to be created")
+	}
+	if first != second {
+		t.Fatalf("expected same turn ID, got %q then %q", first, second)
+	}
+
+	turns, err := repo.ListTurnsBySession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("ListTurnsBySession: %v", err)
+	}
+	if len(turns) != 1 {
+		t.Fatalf("expected 1 turn, got %d", len(turns))
+	}
+}
+
+// TestCompleteTurnClosesUntrackedDBTurn covers the user-cancel zombie path:
+// completeTurnForSession previously bailed when activeTurns was empty, leaving
+// the DB row open (e.g. after a backend restart wiped activeTurns, or after
+// the dual-creation drift left a turn the orchestrator never tracked).
+func TestCompleteTurnClosesUntrackedDBTurn(t *testing.T) {
+	svc, repo := newTurnLifecycleTestService(t)
+	ctx := context.Background()
+
+	// Simulate an open turn the orchestrator never tracked.
+	if _, err := svc.turnService.StartTurn(ctx, "session1"); err != nil {
+		t.Fatalf("seed turn: %v", err)
+	}
+	if open := openTurnCount(t, repo, "session1"); open != 1 {
+		t.Fatalf("expected 1 open turn before complete, got %d", open)
+	}
+
+	svc.completeTurnForSession(ctx, "session1")
+
+	if open := openTurnCount(t, repo, "session1"); open != 0 {
+		t.Fatalf("expected 0 open turns after complete, got %d", open)
+	}
+}
+
+// TestCompleteTurnMopsUpMultipleZombies verifies the loop that cleans up
+// pre-existing zombies from before this fix shipped.
+func TestCompleteTurnMopsUpMultipleZombies(t *testing.T) {
+	svc, repo := newTurnLifecycleTestService(t)
+	ctx := context.Background()
+
+	for i := 0; i < 4; i++ {
+		if _, err := svc.turnService.StartTurn(ctx, "session1"); err != nil {
+			t.Fatalf("seed turn %d: %v", i, err)
+		}
+	}
+	if open := openTurnCount(t, repo, "session1"); open != 4 {
+		t.Fatalf("expected 4 open turns, got %d", open)
+	}
+
+	svc.completeTurnForSession(ctx, "session1")
+
+	if open := openTurnCount(t, repo, "session1"); open != 0 {
+		t.Fatalf("expected 0 open turns after sweep, got %d", open)
+	}
+}
+
+// TestCompleteTurnRespectsIterationCap verifies that the cleanup loop closes at
+// most maxIterations (16) turns per call and that a subsequent call mops up the
+// remainder. Locks in the cap behavior so future tweaks don't accidentally turn
+// the safety bound into a footgun.
+func TestCompleteTurnRespectsIterationCap(t *testing.T) {
+	svc, repo := newTurnLifecycleTestService(t)
+	ctx := context.Background()
+
+	const seeded = 20
+	for i := 0; i < seeded; i++ {
+		if _, err := svc.turnService.StartTurn(ctx, "session1"); err != nil {
+			t.Fatalf("seed turn %d: %v", i, err)
+		}
+	}
+	if open := openTurnCount(t, repo, "session1"); open != seeded {
+		t.Fatalf("expected %d open turns, got %d", seeded, open)
+	}
+
+	svc.completeTurnForSession(ctx, "session1")
+
+	if open := openTurnCount(t, repo, "session1"); open != seeded-16 {
+		t.Fatalf("expected %d open turns after first sweep (cap=16), got %d", seeded-16, open)
+	}
+
+	svc.completeTurnForSession(ctx, "session1")
+
+	if open := openTurnCount(t, repo, "session1"); open != 0 {
+		t.Fatalf("expected 0 open turns after second sweep, got %d", open)
+	}
+}
+
+// TestCompleteTurnIsIdempotent covers the cancel-after-agent-already-completed
+// race: the agent's stream complete event closed the turn, then CancelAgent
+// runs completeTurnForSession again. Should be a no-op, not an error.
+func TestCompleteTurnIsIdempotent(t *testing.T) {
+	svc, repo := newTurnLifecycleTestService(t)
+	ctx := context.Background()
+
+	turnID := svc.startTurnForSession(ctx, "session1")
+	if turnID == "" {
+		t.Fatal("expected turn to be created")
+	}
+
+	svc.completeTurnForSession(ctx, "session1")
+	svc.completeTurnForSession(ctx, "session1") // second call: no-op
+
+	if open := openTurnCount(t, repo, "session1"); open != 0 {
+		t.Fatalf("expected 0 open turns, got %d", open)
+	}
+	turns, err := repo.ListTurnsBySession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("ListTurnsBySession: %v", err)
+	}
+	if len(turns) != 1 {
+		t.Fatalf("expected exactly 1 turn (no phantom), got %d", len(turns))
+	}
+}
@@ -89,20 +89,32 @@ func (s *Service) CompleteTurn(ctx context.Context, turnID string) error {
 }

 // GetActiveTurn returns the currently active (non-completed) turn for a session.
+// Returns (nil, nil) when no turn is active. Other errors (DB failures) are returned as-is.
 func (s *Service) GetActiveTurn(ctx context.Context, sessionID string) (*models.Turn, error) {
-	return s.turns.GetActiveTurnBySessionID(ctx, sessionID)
+	turn, err := s.turns.GetActiveTurnBySessionID(ctx, sessionID)
+	if errors.Is(err, sql.ErrNoRows) {
+		return nil, nil
+	}
+	return turn, err
 }

 // getOrStartTurn returns the active turn for a session, or starts a new one if none exists.
 // This is used to ensure messages always have a valid turn ID.
 func (s *Service) getOrStartTurn(ctx context.Context, sessionID string) (*models.Turn, error) {
-	// First try to get an active turn
-	turn, err := s.turns.GetActiveTurnBySessionID(ctx, sessionID)
-	if err == nil && turn != nil {
+	// Route through GetActiveTurn so the ErrNoRows → (nil, nil) normalization
+	// applies consistently. A real DB read failure is logged before falling
+	// through to StartTurn — same observability contract as the orchestrator's
+	// startTurnForSession adoption path.
+	turn, err := s.GetActiveTurn(ctx, sessionID)
+	if err != nil {
+		s.logger.Warn("failed to look up active turn; will create a new one",
+			zap.String("session_id", sessionID),
+			zap.Error(err))
+	} else if turn != nil {
 		return turn, nil
 	}

-	// No active turn, start a new one
+	// No active turn (or read failed), start a new one
 	return s.StartTurn(ctx, sessionID)
 }