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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user