fix(backend): skip redundant task-state writes when session state is unchanged (#953)
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user