fix(backend): skip redundant task-state writes when session state is unchanged (#953)

This commit is contained in:
Carlos Florêncio
2026-05-19 18:27:16 +01:00
committed by GitHub
parent 2bfdd518b6
commit b73c00f054
3 changed files with 148 additions and 5 deletions
@@ -514,16 +514,47 @@ func isTerminalSessionState(s models.TaskSessionState) bool {
}
func (s *Service) setSessionWaitingForInput(ctx context.Context, taskID, sessionID string, preloadedSession ...*models.TaskSession) {
s.updateTaskSessionState(ctx, taskID, sessionID, models.TaskSessionStateWaitingForInput, "", false, preloadedSession...)
// Resolve session up front so we can skip the redundant task-state write
// when the session was already WAITING_FOR_INPUT. Without this guard, every
// caller (workflow on_turn_complete + handleCompleteStreamEvent + other
// terminal paths) writes tasks.state=REVIEW on every turn even though the
// state hasn't changed, producing duplicate "task moved to REVIEW" logs and
// unnecessary DB churn.
var session *models.TaskSession
if len(preloadedSession) > 0 && preloadedSession[0] != nil {
session = preloadedSession[0]
} else {
var err error
session, err = s.repo.GetTaskSession(ctx, sessionID)
if err != nil {
// Fall back to legacy behavior — still attempt the task-state
// write so a transient lookup failure doesn't drop a needed
// REVIEW transition.
s.updateTaskSessionState(ctx, taskID, sessionID, models.TaskSessionStateWaitingForInput, "", false)
s.writeTaskReviewState(ctx, taskID)
return
}
}
wasAlreadyWaiting := session.State == models.TaskSessionStateWaitingForInput
s.updateTaskSessionState(ctx, taskID, sessionID, models.TaskSessionStateWaitingForInput, "", false, session)
if wasAlreadyWaiting {
return
}
s.writeTaskReviewState(ctx, taskID)
}
func (s *Service) writeTaskReviewState(ctx context.Context, taskID string) {
if err := s.taskRepo.UpdateTaskState(ctx, taskID, v1.TaskStateReview); err != nil {
s.logger.Error("failed to update task state to REVIEW",
zap.String("task_id", taskID),
zap.Error(err))
} else {
s.logger.Info("task moved to REVIEW state",
zap.String("task_id", taskID))
return
}
s.logger.Info("task moved to REVIEW state",
zap.String("task_id", taskID))
}
func (s *Service) setSessionRunning(ctx context.Context, taskID, sessionID string, preloadedSession ...*models.TaskSession) {
@@ -545,8 +576,18 @@ func (s *Service) setSessionRunning(ctx context.Context, taskID, sessionID strin
return
}
// Skip the redundant task-state write when the session is already RUNNING.
// Tool calls fire many events per turn and each was triggering an
// UpdateTaskState(IN_PROGRESS) write that produced no actual state change
// (2,000+ redundant writes observed on long-running turns).
wasAlreadyRunning := session.State == models.TaskSessionStateRunning
s.updateTaskSessionState(ctx, taskID, sessionID, models.TaskSessionStateRunning, "", true, session)
if wasAlreadyRunning {
return
}
// Office tasks do NOT transition to IN_PROGRESS when their agent's
// session enters RUNNING. The user-facing task status (todo /
// in_review / done / blocked) reflects workflow position, not the
@@ -210,3 +210,102 @@ func TestToolEventsWakeSessionAndTaskTogether(t *testing.T) {
})
}
}
// TestSetSessionRunning_NoRedundantTaskWrites locks in the dedup: when the
// session is already RUNNING, setSessionRunning must not re-write tasks.state.
// Without the guard, every tool_call / tool_update fired UpdateTaskState
// (2,000+ redundant writes observed on long-running turns).
func TestSetSessionRunning_NoRedundantTaskWrites(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
session, err := repo.GetTaskSession(ctx, "s1")
require.NoError(t, err)
session.State = models.TaskSessionStateRunning
require.NoError(t, repo.UpdateTaskSession(ctx, session))
taskRepo := newMockTaskRepo()
svc := createTestService(repo, newMockStepGetter(), taskRepo)
for i := 0; i < 5; i++ {
svc.setSessionRunning(ctx, "t1", "s1")
}
require.Equal(t, 0, taskRepo.stateWrites["t1"],
"setSessionRunning must not write tasks.state when session is already RUNNING")
}
// TestSetSessionWaitingForInput_NoRedundantTaskWrites locks in the dedup for
// the WAITING_FOR_INPUT path. Without the guard, both the workflow
// on_turn_complete transition and handleCompleteStreamEvent were writing
// tasks.state=REVIEW back-to-back on every turn.
func TestSetSessionWaitingForInput_NoRedundantTaskWrites(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
session, err := repo.GetTaskSession(ctx, "s1")
require.NoError(t, err)
session.State = models.TaskSessionStateWaitingForInput
require.NoError(t, repo.UpdateTaskSession(ctx, session))
taskRepo := newMockTaskRepo()
svc := createTestService(repo, newMockStepGetter(), taskRepo)
for i := 0; i < 5; i++ {
svc.setSessionWaitingForInput(ctx, "t1", "s1")
}
require.Equal(t, 0, taskRepo.stateWrites["t1"],
"setSessionWaitingForInput must not write tasks.state when session is already WAITING_FOR_INPUT")
}
// TestSetSessionRunning_WritesOnTransition guards against an over-eager dedup
// regression: when the session was NOT already RUNNING, the task state write
// MUST still happen.
func TestSetSessionRunning_WritesOnTransition(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
session, err := repo.GetTaskSession(ctx, "s1")
require.NoError(t, err)
session.State = models.TaskSessionStateWaitingForInput
require.NoError(t, repo.UpdateTaskSession(ctx, session))
taskRepo := newMockTaskRepo()
svc := createTestService(repo, newMockStepGetter(), taskRepo)
svc.setSessionRunning(ctx, "t1", "s1")
require.Equal(t, 1, taskRepo.stateWrites["t1"],
"setSessionRunning must write tasks.state on actual transition")
require.Equal(t, v1.TaskStateInProgress, taskRepo.updatedStates["t1"])
}
// TestSetSessionWaitingForInput_WritesOnTransition is the symmetric counterpart
// to TestSetSessionRunning_WritesOnTransition: when the session is NOT already
// WAITING_FOR_INPUT, setSessionWaitingForInput MUST still fire the task write.
// Without this guard an accidental inversion of wasAlreadyWaiting would silently
// stop tasks from ever reaching REVIEW.
func TestSetSessionWaitingForInput_WritesOnTransition(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
// Seed session in RUNNING state (the normal pre-condition for a turn completing).
session, err := repo.GetTaskSession(ctx, "s1")
require.NoError(t, err)
session.State = models.TaskSessionStateRunning
require.NoError(t, repo.UpdateTaskSession(ctx, session))
taskRepo := newMockTaskRepo()
svc := createTestService(repo, newMockStepGetter(), taskRepo)
svc.setSessionWaitingForInput(ctx, "t1", "s1")
require.Equal(t, 1, taskRepo.stateWrites["t1"],
"setSessionWaitingForInput must write tasks.state on actual transition")
require.Equal(t, v1.TaskStateReview, taskRepo.updatedStates["t1"])
}
@@ -85,13 +85,15 @@ type mockTaskRepo struct {
mu sync.Mutex
tasks map[string]*v1.Task
updatedStates map[string]v1.TaskState
getTaskErr error // if set, GetTask returns this error
stateWrites map[string]int // per-task UpdateTaskState call count for dedup tests
getTaskErr error // if set, GetTask returns this error
}
func newMockTaskRepo() *mockTaskRepo {
return &mockTaskRepo{
tasks: make(map[string]*v1.Task),
updatedStates: make(map[string]v1.TaskState),
stateWrites: make(map[string]int),
}
}
@@ -111,6 +113,7 @@ func (m *mockTaskRepo) UpdateTaskState(_ context.Context, taskID string, state v
m.mu.Lock()
defer m.mu.Unlock()
m.updatedStates[taskID] = state
m.stateWrites[taskID]++
return nil
}