fix(orchestrator): drain queued peer messages on clarification pause + enqueue fast-path (#2837)

Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>
This commit is contained in:
WaleWangPW
2026-08-20 20:02:05 +08:00
committed by GitHub
parent d3f356876a
commit 9820d2a562
12 changed files with 765 additions and 13 deletions
@@ -77,6 +77,14 @@ type ClarificationInputPauser interface {
PauseForClarificationInput(ctx context.Context, sessionID string) (int, error)
}
type clarificationInputPauserWithOptions interface {
PauseForClarificationInputWithOptions(
ctx context.Context,
sessionID string,
options orchestrator.ClarificationPauseOptions,
) (int, error)
}
// MessageCreator creates messages for clarification requests.
type MessageCreator interface {
CreateClarificationRequestMessages(ctx context.Context, taskID, sessionID, pendingID string, questions []clarification.Question, clarificationContext string) ([]string, error)
@@ -10,6 +10,7 @@ import (
"github.com/google/uuid"
"github.com/kandev/kandev/internal/clarification"
"github.com/kandev/kandev/internal/orchestrator"
"github.com/kandev/kandev/internal/task/models"
"github.com/kandev/kandev/internal/task/service"
ws "github.com/kandev/kandev/pkg/websocket"
@@ -146,7 +147,18 @@ func (h *Handlers) handleAskParentQuestion(ctx context.Context, msg *ws.Message)
// lets an answer race the cancellation of the turn that asked it.
h.setSessionWaitingForInput(ctx, target.childTask.ID, target.childSession.ID)
if h.inputPauser != nil {
if _, pauseErr := h.inputPauser.PauseForClarificationInput(context.WithoutCancel(ctx), target.childSession.ID); pauseErr != nil {
pauseCtx := context.WithoutCancel(ctx)
var pauseErr error
if pauser, ok := h.inputPauser.(clarificationInputPauserWithOptions); ok {
_, pauseErr = pauser.PauseForClarificationInputWithOptions(
pauseCtx,
target.childSession.ID,
orchestrator.ClarificationPauseOptions{},
)
} else {
_, pauseErr = h.inputPauser.PauseForClarificationInput(pauseCtx, target.childSession.ID)
}
if pauseErr != nil {
h.logger.Warn("failed to pause autopilot child before parent question delivery",
zap.String("question_id", questionID), zap.String("session_id", target.childSession.ID), zap.Error(pauseErr))
}
@@ -5,6 +5,7 @@ import (
"encoding/json"
"testing"
"github.com/kandev/kandev/internal/orchestrator"
"github.com/kandev/kandev/internal/task/models"
"github.com/kandev/kandev/internal/task/service"
v1 "github.com/kandev/kandev/pkg/api/v1"
@@ -12,6 +13,20 @@ import (
"github.com/stretchr/testify/require"
)
type parentQuestionPausePolicyRecorder struct {
recordingClarificationInputPauser
options []orchestrator.ClarificationPauseOptions
}
func (p *parentQuestionPausePolicyRecorder) PauseForClarificationInputWithOptions(
_ context.Context,
_ string,
options orchestrator.ClarificationPauseOptions,
) (int, error) {
p.options = append(p.options, options)
return p.count, p.err
}
func seedParentQuestionScenario(t *testing.T, svc *service.Service, repo seedRepo) (*models.Task, *models.Task, *models.TaskSession, *models.TaskSession) {
t.Helper()
ctx := context.Background()
@@ -112,6 +127,44 @@ func TestHandleAskParentQuestion_PersistsRoutesAndPauses(t *testing.T) {
require.Contains(t, status.Entries[0].Content, "Which database should I use?")
}
func TestHandleAskParentQuestion_HoldsUnrelatedChildQueue(t *testing.T) {
svc, repo := newTestTaskService(t)
parent, child, _, childSession := seedParentQuestionScenario(t, svc, repo)
h, orch := newMessageTaskHandler(t, svc, repo)
pauser := &parentQuestionPausePolicyRecorder{}
h.inputPauser = pauser
_, err := orch.queue.QueueMessageWithMetadata(
context.Background(), childSession.ID, child.ID, "unrelated child prompt", "", "user-1", false, nil, nil,
)
require.NoError(t, err)
resp, err := h.handleAskParentQuestion(context.Background(), parentQuestionMessage(t, child.ID, childSession.ID))
require.NoError(t, err)
require.Equal(t, ws.MessageTypeResponse, resp.Type)
require.Len(t, pauser.options, 1)
require.False(t, pauser.options[0].DrainQueuedMessages)
require.Len(t, orch.queue.GetStatus(context.Background(), childSession.ID).Entries, 1)
childAfter, err := repo.GetTaskSession(context.Background(), childSession.ID)
require.NoError(t, err)
require.Equal(t, models.TaskSessionStateWaitingForInput, childAfter.State)
questionID := responseField(t, resp, "question_id")
question, err := svc.GetMessage(context.Background(), questionID)
require.NoError(t, err)
require.Equal(t, models.MessageTypeClarificationRequest, question.Type)
require.Equal(t, "pending", question.Metadata[models.MetaKeyParentQuestionStatus])
_ = parent
}
func responseField(t *testing.T, response *ws.Message, field string) string {
t.Helper()
var payload map[string]interface{}
require.NoError(t, json.Unmarshal(response.Payload, &payload))
value, ok := payload[field].(string)
require.True(t, ok)
return value
}
func TestHandleMessageTask_AnswersParentQuestionIdempotently(t *testing.T) {
svc, repo := newTestTaskService(t)
parent, child, _, childSession := seedParentQuestionScenario(t, svc, repo)
@@ -8,6 +8,8 @@ import (
"github.com/kandev/kandev/internal/task/models"
)
const clarificationAgentDisconnectedMetadataKey = "agent_disconnected"
// sessionHasPendingClarification reports whether the session still has durable
// clarification_request rows awaiting user input. Used to fail closed on
// workflow on_turn_complete while the user can still answer.
@@ -25,6 +27,34 @@ func (s *Service) sessionHasPendingClarification(ctx context.Context, sessionID
return len(msgs) > 0
}
// sessionHasLiveClarification reports whether a current-turn clarification
// still owns the agent turn. A detached bundle remains pending and answerable,
// but it no longer owns a live turn, so a parked session may dispatch a newer
// queued prompt. Repository failures fail closed because the caller cannot
// safely distinguish a detached bundle from a live clarification.
func (s *Service) sessionHasLiveClarification(ctx context.Context, sessionID string) bool {
if sessionID == "" {
return false
}
msgs, err := s.repo.FindActiveClarificationMessagesBySessionID(ctx, sessionID)
if err != nil {
s.logger.Warn("failed to check live clarifications; blocking queue drain",
zap.String("session_id", sessionID),
zap.Error(err))
return true
}
for _, msg := range msgs {
if msg == nil || msg.Metadata == nil {
return true
}
detached, ok := msg.Metadata[clarificationAgentDisconnectedMetadataKey].(bool)
if !ok || !detached {
return true
}
}
return false
}
// turnCompleteBlockedByUserInput reports and applies the workflow barrier for
// durable user-input waits. A pending clarification is a platform pause, not a
// step-completion signal, so all turn-complete transition entrypoints fail
@@ -77,6 +77,33 @@ func TestSessionHasPendingClarificationFailsClosedOnRepositoryError(t *testing.T
}
}
func TestSessionHasLiveClarificationIgnoresDetachedBundle(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "task-detached", "session-detached", "step1")
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
now := time.Now().UTC()
requireNoError(t, repo.CreateTurn(ctx, &models.Turn{
ID: "turn-detached", TaskSessionID: "session-detached", TaskID: "task-detached",
StartedAt: now, CreatedAt: now,
}))
requireNoError(t, repo.CreateMessage(ctx, &models.Message{
ID: "clarification-detached", TaskSessionID: "session-detached", TaskID: "task-detached",
TurnID: "turn-detached", AuthorType: models.MessageAuthorAgent,
Type: models.MessageTypeClarificationRequest, Content: "old question", CreatedAt: now,
Metadata: map[string]any{
"pending_id": "pending-detached", "status": "pending", "agent_disconnected": true,
},
}))
if !svc.sessionHasPendingClarification(ctx, "session-detached") {
t.Fatal("detached clarification must remain a pending answer barrier")
}
if svc.sessionHasLiveClarification(ctx, "session-detached") {
t.Fatal("detached-only clarification must not block a successor queue drain")
}
}
func seedPendingClarificationMessage(t *testing.T, repo interface {
CreateTurn(ctx context.Context, turn *models.Turn) error
CreateMessage(ctx context.Context, message *models.Message) error
@@ -773,7 +773,12 @@ func (s *Service) executeQueuedMessageWithReservation(
if reservation == nil {
reservation = s.queuedDispatchReservationForEntry(reservedSessionID, queuedMsg.ID)
}
defer s.clearQueuedDispatchInFlightIfCurrent(reservedSessionID, reservation)
defer func() {
s.clearQueuedDispatchInFlightIfCurrent(reservedSessionID, reservation)
if s.onQueuedMessageExecutionComplete != nil {
s.onQueuedMessageExecutionComplete()
}
}()
lifecyclePrompt := isLifecycleAutomationOrigin(queuedMsg.Metadata["origin"])
claimEntryID, handoffDone := s.claimQueuedMessageHandoff(
@@ -519,12 +519,32 @@ func (s *Service) dispatchClarificationResumeLocked(ctx context.Context, data cl
return nil
}
// ClarificationPauseOptions controls the queue policy for a clarification
// pause. Human no-answer pauses drain an already queued prompt, while an
// autopilot parent question keeps unrelated child prompts parked until the
// parent answers.
type ClarificationPauseOptions struct {
DrainQueuedMessages bool
}
// PauseForClarificationInput converts a no-answer ask_user_question outcome
// into a platform pause. It detaches the pending clarification so a late user
// answer resumes through the event fallback path, then silently cancels the
// active agent turn without evaluating workflow turn-complete actions. It
// returns the number of clarification bundles detached.
func (s *Service) PauseForClarificationInput(ctx context.Context, sessionID string) (int, error) {
return s.PauseForClarificationInputWithOptions(ctx, sessionID, ClarificationPauseOptions{
DrainQueuedMessages: true,
})
}
// PauseForClarificationInputWithOptions is the explicit pause policy used by
// callers whose queue semantics differ from the human clarification timeout.
func (s *Service) PauseForClarificationInputWithOptions(
ctx context.Context,
sessionID string,
options ClarificationPauseOptions,
) (int, error) {
if sessionID == "" {
return 0, fmt.Errorf("session_id is required")
}
@@ -533,22 +553,13 @@ func (s *Service) PauseForClarificationInput(ctx context.Context, sessionID stri
guard := s.lockCancelInFlightGuard(sessionID)
defer guard.release()
guardedPersistenceStartedAt := time.Now()
session, err := s.repo.GetTaskSession(writeCtx, sessionID)
session, expectedTurnID, expectedTurn, err := s.loadClarificationPauseState(writeCtx, sessionID)
if err != nil {
return 0, fmt.Errorf("load session for clarification pause: %w", err)
return 0, err
}
if session == nil {
return 0, nil
}
expectedTurnID := ""
var expectedTurn *string
if s.turnService != nil {
expectedTurnID, err = s.peekActiveTurnID(writeCtx, sessionID)
if err != nil {
return 0, fmt.Errorf("inspect clarification turn before pause: %w", err)
}
expectedTurn = &expectedTurnID
}
hasPendingClarification := s.sessionHasPendingClarification(writeCtx, sessionID)
detached := 0
@@ -596,9 +607,45 @@ func (s *Service) PauseForClarificationInput(ctx context.Context, sessionID stri
} else if err != nil {
return detached, errors.Join(detachErr, err)
}
if !options.DrainQueuedMessages {
return detached, detachErr
}
if detachErr != nil {
// The durable clarification remains the input authority when detachment
// failed. Starting a successor here would make the row look stale and
// allow work to continue without a recoverable question.
return detached, detachErr
}
// Clarification park is a turn boundary, mirroring the
// pre-#677 contract where handleAgentReady always drained after returning
// from a turn. PauseForClarificationInput holds the cancelInFlight guard
// (line 533) for the entire function, so we must use the *Locked* drain
// variant — the public one would try to re-acquire the same non-reentrant
// sync.Mutex and deadlock. The detached bundle remains in the UI for the
// user to answer, but the workflow will no longer block on it for the new
// turn (PR description flags this trade-off for maintainer review).
s.drainQueuedMessageForPromptableSessionLocked(pauseCtx, sessionID)
return detached, detachErr
}
func (s *Service) loadClarificationPauseState(
ctx context.Context,
sessionID string,
) (*models.TaskSession, string, *string, error) {
session, err := s.repo.GetTaskSession(ctx, sessionID)
if err != nil {
return nil, "", nil, fmt.Errorf("load session for clarification pause: %w", err)
}
if session == nil || s.turnService == nil {
return session, "", nil, nil
}
expectedTurnID, err := s.peekActiveTurnID(ctx, sessionID)
if err != nil {
return nil, "", nil, fmt.Errorf("inspect clarification turn before pause: %w", err)
}
return session, expectedTurnID, &expectedTurnID, nil
}
// cancelAgentSilent cancels the agent turn without creating a visible message
// in the chat. Used by clarification recovery so the cancel-and-retry is seamless.
//
@@ -0,0 +1,217 @@
package orchestrator
import (
"context"
"testing"
"time"
"github.com/kandev/kandev/internal/clarification"
"github.com/kandev/kandev/internal/orchestrator/watcher"
"github.com/kandev/kandev/internal/task/models"
wfmodels "github.com/kandev/kandev/internal/workflow/models"
)
// TestPauseForClarificationInput_DrainsPeerQueueAfterPause pins the T1
// contract with the real detaching canceller and a deterministic dispatch
// barrier. Pause must detach the clarification, reserve the queued peer, and
// leave the detached question pending for a late answer.
func TestPauseForClarificationInput_DrainsPeerQueueAfterPause(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateRunning)
seedExecutorRunning(t, repo, "s1", "t1", "exec-1")
setSessionExecID(t, repo, "s1", "exec-1")
seedPendingClarificationMessage(t, repo, "t1", "s1")
agentMgr := &mockAgentManager{
isAgentRunning: true,
repoForExecutionLookup: repo,
}
store := clarification.NewStore(time.Minute)
store.CreateRequest(&clarification.Request{PendingID: "pending-s1", SessionID: "s1"})
canceller := clarification.NewCanceller(store, repo, nil, testLogger())
svc := createEngineService(t, repo, newMockStepGetter(), agentMgr)
svc.SetClarificationCanceller(canceller)
svc.turnService = &repoBackedTurnService{repo: repo}
workerDone := make(chan struct{})
svc.onQueuedMessageExecutionComplete = func() { close(workerDone) }
// Seed a peer message. Distinct QueuedBy to avoid the queue's
// admission-time auto-merge collapsing it.
_, err := svc.messageQueue.QueueMessageWithMetadata(ctx, "s1", "t1", "peer-1", "", "user-1", false, nil, map[string]interface{}{})
if err != nil {
t.Fatalf("queue peer: %v", err)
}
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 1 {
t.Fatalf("pre-pause queue count = %d, want 1", got)
}
detached, err := svc.PauseForClarificationInput(ctx, "s1")
if err != nil {
t.Fatalf("PauseForClarificationInput: %v", err)
}
if detached != 1 {
t.Fatalf("expected one detached bundle, got %d", detached)
}
select {
case <-workerDone:
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for queued peer worker")
}
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 0 {
t.Fatalf("post-pause queue count = %d, want 0 after reservation", got)
}
clarificationMessage, err := repo.GetMessage(ctx, "clarification-s1")
if err != nil {
t.Fatalf("load detached clarification: %v", err)
}
if clarificationMessage.Metadata["status"] != "pending" || clarificationMessage.Metadata["agent_disconnected"] != true {
t.Fatalf("detached clarification metadata = %#v, want pending and disconnected", clarificationMessage.Metadata)
}
}
// TestPauseForClarificationInput_SameTurnBarrierUnchanged pins the
// contract that BEFORE T1 fires (i.e., on a turn that's still in the
// same turn the clarification was asked), the workflow barrier still
// blocks. The detached bundle's turn_id stays equal to the current
// turn until the drain starts a new turn. The barrier must hold for
// the current turn regardless of T1.
func TestPauseForClarificationInput_SameTurnBarrierUnchanged(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
session, err := repo.GetTaskSession(ctx, "s1")
if err != nil {
t.Fatalf("get session: %v", err)
}
session.State = models.TaskSessionStateRunning
if err := repo.UpdateTaskSession(ctx, session); err != nil {
t.Fatalf("update session state: %v", err)
}
seedExecutorRunning(t, repo, "s1", "t1", "exec-1")
setSessionExecID(t, repo, "s1", "exec-1")
seedPendingClarificationMessage(t, repo, "t1", "s1")
stepGetter := newMockStepGetter()
stepGetter.steps["step1"] = &wfmodels.WorkflowStep{
ID: "step1", WorkflowID: "wf1", Name: "Plan", Position: 0,
Events: wfmodels.StepEvents{
OnTurnComplete: []wfmodels.OnTurnCompleteAction{
{Type: wfmodels.OnTurnCompleteMoveToNext},
},
},
}
stepGetter.steps["step2"] = &wfmodels.WorkflowStep{
ID: "step2", WorkflowID: "wf1", Name: "Implement", Position: 1,
}
agentMgr := &mockAgentManager{repoForExecutionLookup: repo}
svc := createEngineService(t, repo, stepGetter, agentMgr)
svc.turnService = &repoBackedTurnService{repo: repo}
// Manually call handleAgentReady without ever going through
// PauseForClarificationInput — the same-turn barrier must hold.
svc.handleAgentReady(ctx, watcher.AgentEventData{TaskID: "t1", SessionID: "s1"})
task, err := repo.GetTask(ctx, "t1")
if err != nil {
t.Fatalf("get task: %v", err)
}
if task.WorkflowStepID != "step1" {
t.Fatalf("same-turn barrier broken: workflow step = %q, want step1", task.WorkflowStepID)
}
}
// TestPauseForClarificationInput_NewTurnAfterDrainAdvancesWorkflow verifies
// that the dispatched successor can complete its turn and advance the
// workflow while the detached clarification remains answerable.
func TestPauseForClarificationInput_NewTurnAfterDrainAdvancesWorkflow(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
seedExecutorRunning(t, repo, "s1", "t1", "exec-1")
setSessionExecID(t, repo, "s1", "exec-1")
seedPendingClarificationMessage(t, repo, "t1", "s1")
stepGetter := newMockStepGetter()
stepGetter.steps["step1"] = &wfmodels.WorkflowStep{
ID: "step1", WorkflowID: "wf1", Name: "Plan", Position: 0,
Events: wfmodels.StepEvents{OnTurnComplete: []wfmodels.OnTurnCompleteAction{{Type: wfmodels.OnTurnCompleteMoveToNext}}},
}
stepGetter.steps["step2"] = &wfmodels.WorkflowStep{ID: "step2", WorkflowID: "wf1", Name: "Implement", Position: 1}
agentMgr := &mockAgentManager{isAgentRunning: true, repoForExecutionLookup: repo}
store := clarification.NewStore(time.Minute)
store.CreateRequest(&clarification.Request{PendingID: "pending-s1", SessionID: "s1"})
svc := createEngineService(t, repo, stepGetter, agentMgr)
svc.SetClarificationCanceller(clarification.NewCanceller(store, repo, nil, testLogger()))
svc.turnService = &repoBackedTurnService{repo: repo}
workerDone := make(chan struct{})
svc.onQueuedMessageExecutionComplete = func() { close(workerDone) }
_, err := svc.messageQueue.QueueMessageWithMetadata(ctx, "s1", "t1", "peer-1", "", "user-1", false, nil, map[string]interface{}{})
if err != nil {
t.Fatalf("queue peer: %v", err)
}
if _, err := svc.PauseForClarificationInput(ctx, "s1"); err != nil {
t.Fatalf("PauseForClarificationInput: %v", err)
}
select {
case <-workerDone:
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for successor worker")
}
turn, err := repo.GetActiveTurnBySessionID(ctx, "s1")
if err != nil || turn == nil {
t.Fatalf("successor turn = %#v, err=%v", turn, err)
}
svc.handleAgentReady(ctx, watcher.AgentEventData{TaskID: "t1", SessionID: "s1", AgentExecutionID: "exec-1"})
task, err := repo.GetTask(ctx, "t1")
if err != nil {
t.Fatalf("load task after successor completion: %v", err)
}
if task.WorkflowStepID != "step2" {
t.Fatalf("successor completion did not advance workflow: step=%q, want step2", task.WorkflowStepID)
}
clarificationMessage, err := repo.GetMessage(ctx, "clarification-s1")
if err != nil {
t.Fatalf("load detached clarification: %v", err)
}
if clarificationMessage.Metadata["status"] != "pending" || clarificationMessage.Metadata["agent_disconnected"] != true {
t.Fatalf("detached clarification metadata = %#v, want pending and disconnected", clarificationMessage.Metadata)
}
}
// TestPauseForClarificationInput_CancelledRunSkipsDrain pins the safety
// contract: when PauseForClarificationInput early-returns because the
// turn is terminal (e.g., already Cancelled), T1's drain must not run.
// The terminal early-return path (line 566-568) is before the
// cancelAgentSilentExpectedWithGuard call and therefore before line 597
// drain.
func TestPauseForClarificationInput_CancelledRunSkipsDrain(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateCancelled)
seedExecutorRunning(t, repo, "s1", "t1", "exec-1")
setSessionExecID(t, repo, "s1", "exec-1")
seedPendingClarificationMessage(t, repo, "t1", "s1")
agentMgr := &mockAgentManager{repoForExecutionLookup: repo}
svc := createEngineService(t, repo, newMockStepGetter(), agentMgr)
svc.turnService = &repoBackedTurnService{repo: repo}
_, err := svc.messageQueue.QueueMessageWithMetadata(ctx, "s1", "t1", "peer-1", "user-4", "user-4", false, nil, map[string]interface{}{})
if err != nil {
t.Fatalf("queue peer: %v", err)
}
detached, err := svc.PauseForClarificationInput(ctx, "s1")
if err != nil {
t.Fatalf("PauseForClarificationInput: %v", err)
}
if detached != 0 {
t.Fatalf("terminal session should not detach, got %d", detached)
}
// Queue must still hold the peer message — drain was bypassed.
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 1 {
t.Fatalf("terminal early-return drained queue (count=%d), want 1", got)
}
}
@@ -930,6 +930,11 @@ func TestPauseForClarificationInput_ContinuesHardPauseAfterDetachFailure(t *test
svc := createEngineService(t, repo, newMockStepGetter(), agentMgr)
svc.SetClarificationCanceller(&failingDetachClarificationCanceller{err: wantErr})
svc.turnService = &repoBackedTurnService{repo: repo}
if _, err := svc.messageQueue.QueueMessageWithMetadata(
ctx, "s1", "t1", "peer-during-detach-failure", "user-1", "user-1", false, nil, map[string]interface{}{},
); err != nil {
t.Fatalf("queue peer: %v", err)
}
detached, err := svc.PauseForClarificationInput(ctx, "s1")
if detached != 0 || !errors.Is(err, wantErr) {
@@ -945,6 +950,9 @@ func TestPauseForClarificationInput_ContinuesHardPauseAfterDetachFailure(t *test
if session.State != models.TaskSessionStateWaitingForInput {
t.Fatalf("expected session waiting for input after detach failure, got %q", session.State)
}
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 1 {
t.Fatalf("detach failure drained queued peer message: count=%d, want 1", got)
}
}
func TestPauseForClarificationInput_RetriesHardPauseAfterSuccessfulDetach(t *testing.T) {
@@ -2474,6 +2474,7 @@ const (
queueDrainSkipped queueDrainOutcome = iota
queueDrainDispatched
queueDrainPaused
queueDrainTaskAdmissionReadFailed
)
func (s *Service) drainQueuedMessageForPromptableSession(ctx context.Context, sessionID string) bool {
@@ -2505,6 +2506,44 @@ func (s *Service) drainQueuedMessageForPromptableSessionOutcome(ctx context.Cont
return s.drainQueuedMessageForPromptableSessionLockedOutcome(ctx, sessionID)
}
// drainQueuedMessageForPromptableSessionWithTaskAdmission is the enqueue-side
// drain. It rechecks WIP admission after taking the session guard and directly
// before reserving the queue head, so a stale pre-enqueue task snapshot cannot
// bypass a workflow capacity wait.
func (s *Service) drainQueuedMessageForPromptableSessionWithTaskAdmission(
ctx context.Context,
taskID, sessionID string,
) queueDrainOutcome {
lock, release := s.acquireCancelInFlightGuard(sessionID)
defer release()
lock.Lock()
defer lock.Unlock()
if s.isCancelInFlight(sessionID) {
return queueDrainSkipped
}
session, err := s.repo.GetTaskSession(ctx, sessionID)
if err != nil {
s.logger.Warn("failed to reload session before admission-gated drain",
zap.String("session_id", sessionID), zap.Error(err))
return queueDrainSkipped
}
if session == nil {
return queueDrainSkipped
}
if err := s.checkSessionPromptable(session.TaskID, sessionID, session.State); err != nil {
return queueDrainSkipped
}
// QueueUserPrompt can race a live clarification between its enqueue-side
// checks and this guarded drain. Re-read the clarification ownership while
// the guard is held so a parent question or human clarification still owns
// the turn when the queue head is reserved. Detached bundles are allowed by
// sessionHasLiveClarification and can therefore make progress here.
if s.sessionHasLiveClarification(ctx, sessionID) {
return queueDrainSkipped
}
return s.drainQueuedMessageForPromptableSessionLockedWithTaskAdmission(ctx, taskID, sessionID)
}
// drainQueuedMessageForPromptableSessionLocked takes the next queued
// message and dispatches it for execution. Callers must ensure the session
// is ready for input first, AND must already hold sessionID's
@@ -2521,10 +2560,28 @@ func (s *Service) drainQueuedMessageForPromptableSessionLocked(ctx context.Conte
}
func (s *Service) drainQueuedMessageForPromptableSessionLockedOutcome(ctx context.Context, sessionID string) queueDrainOutcome {
return s.drainQueuedMessageForPromptableSessionLockedWithTaskAdmission(ctx, "", sessionID)
}
func (s *Service) drainQueuedMessageForPromptableSessionLockedWithTaskAdmission(
ctx context.Context,
taskID, sessionID string,
) queueDrainOutcome {
if s.messageQueue == nil || s.isCancelInFlight(sessionID) ||
s.isQueuedDispatchInFlight(sessionID) || s.isSteerInFlight(sessionID) {
return queueDrainSkipped
}
if taskID != "" {
task, err := s.repo.GetTask(ctx, taskID)
if err != nil {
s.logger.Warn("failed to reload task before admission-gated queue reservation",
zap.String("task_id", taskID), zap.String("session_id", sessionID), zap.Error(err))
return queueDrainTaskAdmissionReadFailed
}
if task == nil || (!task.WIPAdmitted && task.QueuedForStepID != "") {
return queueDrainSkipped
}
}
queuedMsg, ok, autoRun := s.messageQueue.ReserveQueuedWithAutoRun(ctx, sessionID)
if !autoRun {
return queueDrainPaused
@@ -0,0 +1,221 @@
package orchestrator
import (
"context"
"errors"
"sync/atomic"
"testing"
"time"
"github.com/kandev/kandev/internal/task/models"
sqliterepo "github.com/kandev/kandev/internal/task/repository/sqlite"
v1 "github.com/kandev/kandev/pkg/api/v1"
)
type fastPathTaskReadRetryRepo struct {
*sqliterepo.Repository
failNext atomic.Bool
}
func (r *fastPathTaskReadRetryRepo) GetTask(ctx context.Context, taskID string) (*models.Task, error) {
if r.failNext.CompareAndSwap(true, false) {
return nil, errors.New("transient task read failure")
}
return r.Repository.GetTask(ctx, taskID)
}
// TestQueueUserPrompt_T2FastPathDrainsPromptableSession pins the
// T2 contract: a user message admitted via QueueUserPrompt triggers
// a synchronous fast-path drain when the session is ready for input,
// no in-flight cancellation/queue/steer holds the guard, and the task
// is not in WIP-wait. The drain uses the existing public helper
// drainQueuedMessageForPromptableSession.
//
// The unit test only verifies the T2 call site is reached (drain's
// ReserveQueued removes the head from the queue, so count drops to
// 0). The downstream dispatch (promptTask) requires a working mock
// agent; full e2e is covered by integration tests.
func TestQueueUserPrompt_T2FastPathDrainsPromptableSession(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 0 {
t.Fatalf("pre-condition: queue count = %d, want 0", got)
}
if err := svc.QueueUserPrompt(ctx, "t1", "s1", "hello", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
// drainQueuedMessageForPromptableSession reserves the head before
// dispatching. The reservation removes the entry from the queue,
// so the count drops to 0 once T2's fast-path drain is reached.
// The downstream dispatch may fail (mock agent can't resume) but
// the count==0 invariant pins the T2 call site.
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 0 {
t.Fatalf("post-enqueue queue count = %d, want 0 (T2 fast-path did not drain)", got)
}
}
// TestQueueUserPrompt_T2SkipsFastPathOnPendingClarification pins the
// live-clarification contract: a session with an active clarification
// request must NOT fast-path drain on enqueue. The drain is deferred to
// the clarification outcome or the next turn boundary.
func TestQueueUserPrompt_T2SkipsFastPathOnPendingClarification(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
seedPendingClarificationMessage(t, repo, "t1", "s1")
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
if err := svc.QueueUserPrompt(ctx, "t1", "s1", "during-clarification", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
// T2 must NOT drain — the user has not yet answered the
// clarification. The queue grows; the next turn-end will drain.
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 1 {
t.Fatalf("post-enqueue queue count = %d, want 1 (T2 must defer to turn-end on pending clarification)", got)
}
}
func TestQueueUserPrompt_T2DrainsAfterClarificationDetachedWithEmptyQueue(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
seedPendingClarificationMessage(t, repo, "t1", "s1")
message, err := repo.GetMessage(ctx, "clarification-s1")
if err != nil {
t.Fatalf("load clarification: %v", err)
}
message.Metadata["agent_disconnected"] = true
if err := repo.UpdateMessage(ctx, message); err != nil {
t.Fatalf("mark clarification detached: %v", err)
}
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
if err := svc.QueueUserPrompt(ctx, "t1", "s1", "after-detach", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 0 {
t.Fatalf("detached-only clarification stranded successor queue: count=%d, want 0", got)
}
}
func TestQueueUserPrompt_T2RetriesTaskAdmissionReadAfterPromotion(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
retryRepo := &fastPathTaskReadRetryRepo{Repository: repo}
retryRepo.failNext.Store(true)
svc.repo = retryRepo
if err := svc.QueueUserPrompt(ctx, "t1", "s1", "after-promotion", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 0 {
t.Fatalf("transient task read stranded queue entry: count=%d, want 0", got)
}
}
// TestQueueUserPrompt_T2SkipsFastPathWhenInFlight pins the in-flight
// guard contract: a queue dispatch already in flight must not be
// raced by the fast-path drain. The drain's internal guard check
// bails when isQueuedDispatchInFlight is true.
func TestQueueUserPrompt_T2SkipsFastPathWhenInFlight(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
// Simulate an in-flight dispatch: the gate is set by the dispatcher's
// own markQueuedDispatchInFlight. Set the same internal flag the
// fast-path checks.
svc.markQueuedDispatchInFlightWithSourceLocked("s1", "in-flight-source-id", nil)
if err := svc.QueueUserPrompt(ctx, "t1", "s1", "during-in-flight", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
// T2 must NOT drain — the drain would race the in-flight dispatch.
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 1 {
t.Fatalf("post-enqueue queue count = %d, want 1 (T2 must not race in-flight dispatch)", got)
}
}
// TestQueueUserPrompt_T2SkipsFastPathOnWIPWait pins the WIP admission
// contract: a task that has not been WIP-admitted and is queued for
// admission (QueuedForStepID != "") must NOT fast-path drain. The
// next admission promotion runs the drain via the existing reconcile
// path; promoting here would bypass WIP gating.
func TestQueueUserPrompt_T2SkipsFastPathOnWIPWait(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
// A task in WIP-wait: WIPAdmitted=false, QueuedForStepID="step1".
task := &models.Task{
ID: "t-wip",
WorkflowID: "wf1",
Title: "WIP-wait task",
State: v1.TaskStateInProgress,
QueuedForStepID: "step1",
WIPAdmitted: false,
CreatedAt: time.Now().UTC(),
UpdatedAt: time.Now().UTC(),
}
if err := repo.CreateTask(ctx, task); err != nil {
t.Fatalf("create task: %v", err)
}
// Session in a state that lets the test focus on the WIP guard.
if err := repo.CreateTaskSession(ctx, &models.TaskSession{
ID: "s-wip",
TaskID: "t-wip",
State: models.TaskSessionStateWaitingForInput,
StartedAt: time.Now().UTC(),
UpdatedAt: time.Now().UTC(),
}); err != nil {
t.Fatalf("create session: %v", err)
}
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
if err := svc.QueueUserPrompt(ctx, "t-wip", "s-wip", "during-wip-wait", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
// T2 must NOT drain — the task is in WIP-wait; the admission
// promotion path is the one that drains.
if got := svc.messageQueue.GetStatus(ctx, "s-wip").Count; got != 1 {
t.Fatalf("post-enqueue queue count = %d, want 1 (T2 must not bypass WIP wait)", got)
}
}
// TestQueueUserPrompt_T2DrainsWhenTaskAdmitted pins the positive
// WIP-admitted case: a task with WIPAdmitted=true (or no
// QueuedForStepID) is ready for prompt, and the fast-path drain
// fires. drainQueuedMessageForPromptableSession reserves the head,
// removing it from the queue (count → 0). The dispatcher's own
// promptTask picks up the queue head but the mock agent cannot
// truly resume — that doesn't change the count==0 invariant.
func TestQueueUserPrompt_T2DrainsWhenTaskAdmitted(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
// A WIP-admitted task (or no WIP wait at all).
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
// Mark the task as WIP-admitted (seedTaskAndSession defaults to
// v1.TaskStateInProgress and may not set the WIPAdmitted field).
task, err := repo.GetTask(ctx, "t1")
if err != nil {
t.Fatalf("get task: %v", err)
}
task.WIPAdmitted = true
if err := repo.UpdateTask(ctx, task); err != nil {
t.Fatalf("update task: %v", err)
}
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
if err := svc.QueueUserPrompt(ctx, "t1", "s1", "admitted-task", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
// T2 should drain immediately — the task is admitted and the
// session is ready.
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 0 {
t.Fatalf("post-enqueue queue count = %d, want 0 (T2 fast-path drained admitted task)", got)
}
}
@@ -558,6 +558,11 @@ type Service struct {
// onManualMoveLifecycleStart blocks the admitted manual-move lifecycle in
// package tests so feeder promotion ordering can be asserted.
onManualMoveLifecycleStart func()
// onQueuedMessageExecutionComplete is a package-test hook that fires after
// an asynchronous queued-message worker has released its dispatch
// reservation. It lets tests wait for the full worker lifecycle instead of
// observing only the prompt call.
onQueuedMessageExecutionComplete func()
// queuedMoveLifecycleLocks serializes source-exit work per task. The
// completion marker remains durable so a restart can safely resume work.
queuedMoveLifecycleLocks sync.Map
@@ -3021,9 +3026,71 @@ func (s *Service) QueueUserPrompt(
return fmt.Errorf("queue user prompt: %w", err)
}
s.publishQueueStatusEvent(ctx, sessionID)
// T2: enqueue-side fast-path drain. The user's WIP wait is a
// first-class contract (the dispatcher gates on
// ProcessOnTurnStart's Queued=true return value). Promoting a queued
// prompt to dispatch while the task is still in WIP-wait would
// bypass the admission path; the existing 4 drain triggers
// (agent.ready / agent.boot_ready / processOnEnter / promptTask
// tail) lack this guard too — T2 is defense-in-depth, matching the
// queueFull-bound pattern rather than introducing a new escape.
// Three SQL reads per enqueue (clarification guard, session, and task)
// is small and acceptable for the user-message volume; the load-path hot
// loop is already gated by the guarded clarification check and
// isCancelInFlight / isQueuedDispatchInFlight / isSteerInFlight
// in-flight checks (cheaper than the DB read).
s.tryFastPathDrainAfterEnqueue(ctx, taskID, sessionID)
return nil
}
// tryFastPathDrainAfterEnqueue is the T2 fast-path drain. After QueueUserPrompt
// has admitted a user message into the session's queue, this method
// decides whether the message is safe to dispatch right now. The
// decision is gated by a sequence of cheap-to-expensive checks in
// increasing cost order:
//
// 1. messageQueue is non-nil (cheap pointer check).
// 2. isCancelInFlight / isQueuedDispatchInFlight / isSteerInFlight
// are false (in-memory atomic loads). These guard the next drain
// from racing an in-flight sibling. False negatives are safe: the
// next turn-end re-evaluates.
// 3. drainQueuedMessageForPromptableSessionWithTaskAdmission reloads the
// task after acquiring the guard and immediately before reservation. A
// task waiting in the WIP admission queue (QueuedForStepID != "" &&
// !WIPAdmitted) remains parked.
//
// The first failure aborts the fast-path
// drain silently. Failures leave the queued message in place for the
// next turn-end drain — they never block user input.
//
// The guarded helper rechecks clarification ownership, session promptability,
// and task admission immediately before reservation. The four pre-existing
// drain triggers (agent.ready, agent.boot_ready, processOnEnter, promptTask
// tail) also lack a WIP check — T2 is defense-in-depth that matches the
// existing admission pattern at the queue boundary, rather than introducing a
// different escape path inside the dispatch lifecycle.
func (s *Service) tryFastPathDrainAfterEnqueue(ctx context.Context, taskID, sessionID string) {
if s.messageQueue == nil {
return
}
if s.isCancelInFlight(sessionID) || s.isQueuedDispatchInFlight(sessionID) || s.isSteerInFlight(sessionID) {
return
}
const maxTaskAdmissionReadAttempts = 2
for attempt := 0; attempt < maxTaskAdmissionReadAttempts; attempt++ {
outcome := s.drainQueuedMessageForPromptableSessionWithTaskAdmission(ctx, taskID, sessionID)
if outcome != queueDrainTaskAdmissionReadFailed {
return
}
if attempt+1 == maxTaskAdmissionReadAttempts {
s.logger.Warn("queue fast-path deferred after repeated task admission reads failed",
zap.String("task_id", taskID), zap.String("session_id", sessionID))
return
}
}
}
// GetEventBus returns the event bus
func (s *Service) GetEventBus() bus.EventBus {
return s.eventBus