fix(orchestrator): stop messages sent during a step move from vanishing (#3494)

Co-authored-by: nova28 <17953305+nova28@users.noreply.github.com>
Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>
This commit is contained in:
Henry
2026-09-09 02:53:32 +08:00
committed by GitHub
parent 7435bc2e8b
commit 623ed4e8d5
15 changed files with 806 additions and 33 deletions
@@ -389,6 +389,9 @@ func (s *Service) handleAgentBootReady(ctx context.Context, data watcher.AgentEv
// above is the guarded admission decision; the drain performs its own
// cancellation check and leaves the queue untouched if a new cancellation
// claims the session in this handoff.
if s.isQueuedDispatchInFlight(data.SessionID) {
s.markQueuedDispatchDrainPending(data.SessionID)
}
lock.Unlock()
guardLocked = false
s.drainQueuedMessageForPromptableSession(ctx, data.SessionID)
@@ -818,6 +821,7 @@ func (s *Service) executeQueuedMessageWithReservation(
}
defer func() {
s.clearQueuedDispatchInFlightIfCurrent(reservedSessionID, reservation)
s.drainQueuedDispatchIfPending(reservedSessionID)
if s.onQueuedMessageExecutionComplete != nil {
s.onQueuedMessageExecutionComplete()
}
@@ -888,7 +892,9 @@ func (s *Service) executeQueuedMessageWithReservation(
zap.String("queue_id", queuedMsg.ID))
return
}
s.processOnTurnStartViaEngine(promptCtx, queuedMsg.TaskID, session)
if !turnStartAlreadyProcessed(queuedMsg.Metadata) {
s.processOnTurnStartViaEngine(promptCtx, queuedMsg.TaskID, session)
}
}
// Call promptTask with this entry's ID as a second ownership check. The
@@ -998,10 +1004,13 @@ func (s *Service) handleQueuedMessageExecutionError(
zap.Error(err))
manualRecovery := isManualRecoveryPromptError(err)
// ErrSessionRuntimeUnavailable: the runtime for a just-promoted session has
// not finished launching. Requeue so the drain on agent.boot_ready delivers it.
if lifecyclePrompt || errors.Is(err, errLifecyclePromptClaim) ||
errors.Is(err, errLifecyclePromptMessagePersistence) ||
isSessionBusyError(err) || isTransientPromptError(err) || manualRecovery ||
errors.Is(err, lifecycle.ErrCancelEscalated) || isSessionResetInProgressError(err) {
errors.Is(err, lifecycle.ErrCancelEscalated) || isSessionResetInProgressError(err) ||
errors.Is(err, ErrSessionRuntimeUnavailable) {
if userMessageRecorded {
markQueuedUserMessageRecorded(queuedMsg)
}
@@ -13,6 +13,7 @@ import (
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
"github.com/kandev/kandev/internal/sysprompt"
"github.com/kandev/kandev/internal/task/models"
wfmodels "github.com/kandev/kandev/internal/workflow/models"
v1 "github.com/kandev/kandev/pkg/api/v1"
)
@@ -152,6 +153,73 @@ func TestExecuteQueuedMessage_SkipsUserMessageWhenAlreadyRecorded(t *testing.T)
}
}
// TestExecuteQueuedMessage_SkipsOnTurnStartWhenAlreadyProcessed pins the
// double-fire fix: wsAddMessage's queuePromptIfRuntimeUnavailable path queues
// a prompt after ProcessOnTurnStart already ran synchronously for it, so the
// queued-dispatch path must not fire on_turn_start a second time. Without the
// MetaKeyTurnStartAlreadyProcessed guard, this drain would move the task's
// step twice for one prompt.
func TestExecuteQueuedMessage_SkipsOnTurnStartWhenAlreadyProcessed(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("failed to get session: %v", err)
}
session.State = models.TaskSessionStateWaitingForInput
session.AgentExecutionID = "exec-1"
if err := repo.UpdateTaskSession(ctx, session); err != nil {
t.Fatalf("failed to update session: %v", err)
}
stepGetter := newMockStepGetter()
stepGetter.steps["step1"] = &wfmodels.WorkflowStep{
ID: "step1", WorkflowID: "wf1", Name: "Step 1", Position: 0,
Events: wfmodels.StepEvents{
OnTurnStart: []wfmodels.OnTurnStartAction{
{Type: wfmodels.OnTurnStartMoveToNext},
},
},
}
stepGetter.steps["step2"] = &wfmodels.WorkflowStep{
ID: "step2", WorkflowID: "wf1", Name: "Step 2", Position: 1,
Events: wfmodels.StepEvents{},
}
taskRepo := newMockTaskRepo()
agentMgr := &mockAgentManager{isAgentRunning: true, repoForExecutionLookup: repo}
svc := createTestServiceWithAgent(repo, stepGetter, taskRepo, agentMgr)
svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
seedExecutorRunning(t, repo, "s1", "t1", "exec-1")
queuedMsg := &messagequeue.QueuedMessage{
ID: "q1",
SessionID: "s1",
TaskID: "t1",
Content: "hello",
QueuedBy: "test",
Metadata: map[string]interface{}{
MetaKeyTurnStartAlreadyProcessed: true,
},
}
svc.markQueuedDispatchInFlight("s1", queuedMsg.ID)
svc.executeQueuedMessage("s1", queuedMsg)
task, err := repo.GetTask(ctx, "t1")
if err != nil {
t.Fatalf("failed to get task: %v", err)
}
if task.WorkflowStepID != "step1" {
t.Fatalf("on_turn_start fired a second time: task moved to %q, want it to stay on step1", task.WorkflowStepID)
}
if len(agentMgr.capturedPrompts) != 1 {
t.Fatalf("expected the prompt to still reach PromptAgent, captured=%d", len(agentMgr.capturedPrompts))
}
}
func TestExecuteQueuedMessage_RecordsCIAutomationPromptOnDrain(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
@@ -3,15 +3,87 @@ package orchestrator
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/kandev/kandev/internal/orchestrator/executor"
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
"github.com/kandev/kandev/internal/orchestrator/watcher"
"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 bootReadyRequeueBarrierRepo struct {
*sqliterepo.Repository
lookupStarted chan struct{}
allowLookup chan struct{}
lookupOnce sync.Once
}
func (r *bootReadyRequeueBarrierRepo) GetExecutorRunningBySessionID(
ctx context.Context,
sessionID string,
) (*models.ExecutorRunning, error) {
r.lookupOnce.Do(func() {
close(r.lookupStarted)
<-r.allowLookup
})
return nil, models.ErrExecutorRunningNotFound
}
func TestAgentBootReadyDrainsAfterInFlightRuntimeUnavailableRequeue(t *testing.T) {
ctx := context.Background()
baseRepo := setupTestRepo(t)
seedTaskAndSession(t, baseRepo, "t1", "s1", models.TaskSessionStateWaitingForInput)
repo := &bootReadyRequeueBarrierRepo{
Repository: baseRepo,
lookupStarted: make(chan struct{}),
allowLookup: make(chan struct{}),
}
agentMgr := &mockAgentManager{repoForExecutionLookup: baseRepo}
svc := createTestServiceWithAgent(baseRepo, newMockStepGetter(), newMockTaskRepo(), agentMgr)
svc.repo = repo
svc.executor = executor.NewExecutor(agentMgr, baseRepo, testLogger(), executor.ExecutorConfig{})
if _, err := svc.messageQueue.QueueMessage(ctx, "s1", "t1", "queued prompt", "", messagequeue.QueuedByUser, false, nil); err != nil {
t.Fatalf("queue prompt: %v", err)
}
queued, ok := svc.messageQueue.ReserveQueued(ctx, "s1")
if !ok {
t.Fatal("reserve queued prompt")
}
reservation := svc.markQueuedDispatchInFlightWithSource("s1", queued.ID, queued)
var completions atomic.Int32
secondExecutionDone := make(chan struct{})
svc.onQueuedMessageExecutionComplete = func() {
if completions.Add(1) == 2 {
close(secondExecutionDone)
}
}
go svc.executeQueuedMessageWithReservation("s1", queued, reservation)
select {
case <-repo.lookupStarted:
case <-time.After(time.Second):
t.Fatal("timed out waiting for runtime lookup barrier")
}
svc.handleAgentBootReady(ctx, watcher.AgentEventData{TaskID: "t1", SessionID: "s1"})
close(repo.allowLookup)
select {
case <-secondExecutionDone:
case <-time.After(2 * time.Second):
t.Fatalf("boot-ready did not retry after in-flight requeue, executions=%d", completions.Load())
}
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 1 {
t.Fatalf("runtime-unavailable retry queue count = %d, want 1", got)
}
}
type turnStartSignalService struct {
TurnService
started chan<- struct{}
@@ -1066,6 +1066,29 @@ func (s *Service) updateTaskSessionStateWithHook(
onChanged func(),
preloadedSession ...*models.TaskSession,
) (*models.TaskSession, bool) {
if isTerminalSessionState(nextState) && s.messageQueue != nil {
heldSessionID, _ := ctx.Value(sessionPromptAdmissionContextKey{}).(string)
if heldSessionID != sessionID {
var updated *models.TaskSession
var changed bool
err := s.withSessionPromptAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
updated, changed = s.updateTaskSessionStateWithHook(
admittedCtx,
taskID,
sessionID,
nextState,
errorMessage,
allowWakeFromWaiting,
onChanged,
)
return nil
})
if err != nil {
return nil, false
}
return updated, changed
}
}
var session *models.TaskSession
if len(preloadedSession) > 0 && preloadedSession[0] != nil {
session = preloadedSession[0]
@@ -1229,6 +1252,26 @@ func (s *Service) transitionTaskSessionState(
errorMessage string,
onChanged func(),
) (bool, models.TaskSessionState, error) {
if isTerminalSessionState(nextState) && s.messageQueue != nil {
heldSessionID, _ := ctx.Value(sessionPromptAdmissionContextKey{}).(string)
if heldSessionID != sessionID {
var changed bool
var finalState models.TaskSessionState
var err error
err = s.withSessionPromptAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
changed, finalState, err = s.transitionTaskSessionState(
admittedCtx,
taskID,
sessionID,
nextState,
errorMessage,
onChanged,
)
return err
})
return changed, finalState, err
}
}
session, err := s.repo.GetTaskSession(ctx, sessionID)
if err != nil {
return false, "", fmt.Errorf("get session before state transition: %w", err)
@@ -3909,6 +3909,20 @@ func (s *Service) autoStartPassthroughPrompt(
// failed transiently and the queue was later drained via boot_ready.
const metaKeyUserMessageRecorded = "user_message_recorded"
// MetaKeyTurnStartAlreadyProcessed marks a queued prompt whose on_turn_start
// hook already ran synchronously before queuing (see
// task/handlers.queuePromptIfRuntimeUnavailable, which queues after
// wsAddMessage already called ProcessOnTurnStart for the same user message).
// executeQueuedMessageWithReservation reads this flag to skip its own
// processOnTurnStartViaEngine call and avoid firing on_turn_start twice for
// one prompt.
const MetaKeyTurnStartAlreadyProcessed = "turn_start_already_processed"
func turnStartAlreadyProcessed(metadata map[string]interface{}) bool {
processed, _ := metadata[MetaKeyTurnStartAlreadyProcessed].(bool)
return processed
}
type workflowMessageOrigin struct {
StepID string
StepName string
@@ -398,7 +398,10 @@ func (s *Service) executeSendNowClaimWithContext(
if reservation == nil {
reservation = s.queuedDispatchReservationForEntry(sessionID, claim.Dispatch.ID)
}
defer s.clearQueuedDispatchInFlightIfCurrent(sessionID, reservation)
defer func() {
s.clearQueuedDispatchInFlightIfCurrent(sessionID, reservation)
s.drainQueuedDispatchIfPending(sessionID)
}()
restore := func() {
if err := s.restoreSendNowClaimWithRetry(ctx, claim); err != nil {
@@ -472,7 +475,9 @@ func (s *Service) promptSendNowClaim(ctx context.Context, claim *messagequeue.Se
// rejects the replacement. The ordinary FIFO handoff uses the same
// ordering: workflow admission precedes executor prompt acceptance, and
// the restored claim is retried through the normal ready path.
s.processOnTurnStartViaEngine(ctx, claim.Dispatch.TaskID, session)
if !turnStartAlreadyProcessed(claim.Dispatch.Metadata) {
s.processOnTurnStartViaEngine(ctx, claim.Dispatch.TaskID, session)
}
}
_, err := s.promptTask(ctx, claim.Dispatch.TaskID, sessionID, promptContent, claim.Dispatch.Model,
@@ -13,6 +13,7 @@ import (
"github.com/kandev/kandev/internal/orchestrator/executor"
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
"github.com/kandev/kandev/internal/task/models"
wfmodels "github.com/kandev/kandev/internal/workflow/models"
v1 "github.com/kandev/kandev/pkg/api/v1"
)
@@ -185,6 +186,69 @@ func TestExecuteSendNowClaimRestorePreservesRecordedSources(t *testing.T) {
}
}
func TestPromptSendNowClaimSkipsOnTurnStartWhenAlreadyProcessed(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "task-1", "session-1", "step-1")
seedExecutorRunning(t, repo, "session-1", "task-1", "exec-1")
session, err := repo.GetTaskSession(ctx, "session-1")
if err != nil {
t.Fatalf("get session: %v", err)
}
session.State = models.TaskSessionStateWaitingForInput
session.AgentExecutionID = "exec-1"
if err := repo.UpdateTaskSession(ctx, session); err != nil {
t.Fatalf("set session waiting: %v", err)
}
stepGetter := newMockStepGetter()
stepGetter.steps["step-1"] = &wfmodels.WorkflowStep{
ID: "step-1", WorkflowID: "wf-1", Name: "Step 1", Position: 0,
Events: wfmodels.StepEvents{
OnTurnStart: []wfmodels.OnTurnStartAction{
{Type: wfmodels.OnTurnStartMoveToNext},
},
},
}
stepGetter.steps["step-2"] = &wfmodels.WorkflowStep{
ID: "step-2", WorkflowID: "wf-1", Name: "Step 2", Position: 1,
}
agentMgr := &mockAgentManager{isAgentRunning: true, repoForExecutionLookup: repo}
svc := createTestServiceWithAgent(repo, stepGetter, newMockTaskRepo(), agentMgr)
svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
claim := &messagequeue.SendNowClaim{
Dispatch: messagequeue.QueuedMessage{
ID: "q-1",
SessionID: "session-1",
TaskID: "task-1",
Content: "already admitted",
Metadata: map[string]interface{}{
MetaKeyTurnStartAlreadyProcessed: true,
},
},
}
reservation := svc.markQueuedDispatchInFlight("session-1", claim.Dispatch.ID)
if tracked, err := svc.claimQueuedDispatchForExecution("session-1", claim.Dispatch.ID, reservation); err != nil || !tracked {
t.Fatalf("claim send-now dispatch: tracked=%v err=%v", tracked, err)
}
if err := svc.promptSendNowClaim(ctx, claim); err != nil {
t.Fatalf("prompt send-now claim: %v", err)
}
task, err := repo.GetTask(ctx, "task-1")
if err != nil {
t.Fatalf("get task: %v", err)
}
if task.WorkflowStepID != "step-1" {
t.Fatalf("on_turn_start fired a second time: task moved to %q, want it to stay on step-1", task.WorkflowStepID)
}
if len(agentMgr.capturedPrompts) != 1 {
t.Fatalf("expected the prompt to reach PromptAgent once, captured=%d", len(agentMgr.capturedPrompts))
}
}
func TestSendQueuedNowMissingEntryPreservesAutoRunOff(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
@@ -24,6 +24,21 @@ func (r *fastPathTaskReadRetryRepo) GetTask(ctx context.Context, taskID string)
return r.Repository.GetTask(ctx, taskID)
}
func TestQueueUserPromptRejectsTerminalSession(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateCompleted)
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
err := svc.QueueUserPrompt(ctx, "t1", "s1", "late prompt", "", false, nil, nil, true)
if !errors.Is(err, ErrSessionNotPromptable) {
t.Fatalf("QueueUserPrompt error = %v, want ErrSessionNotPromptable", err)
}
if got := svc.messageQueue.GetStatus(ctx, "s1").Count; got != 0 {
t.Fatalf("terminal session queue count = %d, want 0", got)
}
}
// 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,
@@ -0,0 +1,196 @@
package orchestrator
import (
"context"
"errors"
"testing"
"time"
"github.com/kandev/kandev/internal/orchestrator/executor"
"github.com/kandev/kandev/internal/task/models"
sqliterepo "github.com/kandev/kandev/internal/task/repository/sqlite"
)
// TestQueueUserPrompt_RuntimeUnavailableSurvivesFastPathRedispatch is the
// regression test for Review round 1's F1 finding: QueueUserPrompt's T2
// fast-path drain (tryFastPathDrainAfterEnqueue) reloads the session with
// only checkSessionPromptable, which cannot tell "promptable" apart from
// "promptable but its runtime has not finished launching". So the fast path
// takes the entry and re-dispatches it through promptTask microseconds
// after it was queued — with the runtime exactly as unavailable as it was
// the first time.
//
// This drives the real Service.QueueUserPrompt -> tryFastPathDrainAfterEnqueue
// -> executeQueuedMessageWithReservation -> promptTask -> handleQueuedMessageExecutionError
// chain end to end (no fakes), the way F3 requires: a fake orchestrator that
// only counts QueueUserPrompt calls cannot see this loop at all, because the
// loop happens entirely inside the real Service after the initial enqueue
// returns.
func TestQueueUserPrompt_RuntimeUnavailableSurvivesFastPathRedispatch(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
// WAITING_FOR_INPUT with no executors_running row and no in-memory
// execution reproduces the exact production shape: a session promoted
// by a workflow step move whose runtime has not launched yet.
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
svc.executor = executor.NewExecutor(&mockAgentManager{repoForExecutionLookup: repo}, repo, testLogger(), executor.ExecutorConfig{})
workerDone := make(chan struct{})
svc.onQueuedMessageExecutionComplete = func() { close(workerDone) }
if err := svc.QueueUserPrompt(ctx, "t1", "s1", "hello", "", false, nil, map[string]interface{}{}, true); err != nil {
t.Fatalf("QueueUserPrompt: %v", err)
}
select {
case <-workerDone:
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for the fast-path redispatch to settle")
}
entries, _, err := svc.messageQueue.SnapshotSession(ctx, "s1")
if err != nil {
t.Fatalf("snapshot queue: %v", err)
}
if len(entries) != 1 {
t.Fatalf("message lost: fast-path redispatch failed with ErrSessionRuntimeUnavailable and "+
"was not requeued (queue has %d entries, want 1)", len(entries))
}
if entries[0].Content != "hello" {
t.Fatalf("surviving entry content = %q, want %q", entries[0].Content, "hello")
}
}
// TestPromptTask_SessionAwaitingLaunchClassifiedRuntimeUnavailable pins the
// positive case for the narrowed sentinel (Review round 1's F2 finding): the
// exact "prepared but not launched yet" shape (WAITING_FOR_INPUT, no
// executors_running row) must still classify as ErrSessionRuntimeUnavailable
// so the caller queues instead of surfacing an error.
func TestPromptTask_SessionAwaitingLaunchClassifiedRuntimeUnavailable(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
svc.executor = executor.NewExecutor(&mockAgentManager{repoForExecutionLookup: repo}, repo, testLogger(), executor.ExecutorConfig{})
_, err := svc.PromptTask(ctx, "t1", "s1", "hello", "", false, nil, false)
if err == nil {
t.Fatal("expected an error: no executors_running row exists for this session")
}
if !errors.Is(err, ErrSessionRuntimeUnavailable) {
t.Fatalf("expected ErrSessionRuntimeUnavailable for the not-yet-launched shape, got: %v", err)
}
}
// TestPromptTask_ExhaustedColdResumeNotClassifiedRuntimeUnavailable pins the
// negative case for Review round 1's F2 finding: a session that has an
// executors_running row (so a resume attempt actually runs) and never
// becomes prompt-ready is a genuine launch failure, not the "hasn't started
// yet" case ErrSessionRuntimeUnavailable exists for. It must keep surfacing
// as a plain, visible error so it isn't silently queued forever with no
// error row and no log — the exact regression F2 flagged.
func TestPromptTask_ExhaustedColdResumeNotClassifiedRuntimeUnavailable(t *testing.T) {
oldReadyTimeout := agentPromptReadyTimeout
oldReadyInterval := agentPromptReadyInterval
agentPromptReadyTimeout = 20 * time.Millisecond
agentPromptReadyInterval = time.Millisecond
t.Cleanup(func() {
agentPromptReadyTimeout = oldReadyTimeout
agentPromptReadyInterval = oldReadyInterval
})
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
session, err := repo.GetTaskSession(ctx, "s1")
if err != nil {
t.Fatalf("load session: %v", err)
}
session.AgentProfileID = "profile1"
if err := repo.UpdateTaskSession(ctx, session); err != nil {
t.Fatalf("update session: %v", err)
}
now := time.Now().UTC()
if err := repo.UpsertExecutorRunning(ctx, &models.ExecutorRunning{
ID: "er1",
SessionID: "s1",
TaskID: "t1",
AgentExecutionID: "exec-wedged",
Status: "running",
Resumable: true,
ResumeToken: "resume-token-123",
CreatedAt: now,
UpdatedAt: now,
}); err != nil {
t.Fatalf("seed executor running: %v", err)
}
agentMgr := &mockAgentManager{
repoForExecutionLookup: repo,
isAgentRunningFn: func(context.Context, string) bool { return false },
isAgentReadyFn: func(context.Context, string) bool { return false },
stopAgentWithReasonFunc: func(context.Context, string, string, bool) error {
return nil
},
launchAgentFunc: func(_ context.Context, req *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) {
go func(sessID string) {
sess, gErr := repo.GetTaskSession(ctx, sessID)
if gErr == nil && sess != nil {
sess.State = models.TaskSessionStateWaitingForInput
sess.UpdatedAt = time.Now().UTC()
_ = repo.UpdateTaskSession(ctx, sess)
}
}(req.SessionID)
return &executor.LaunchAgentResponse{AgentExecutionID: "exec-still-wedged"}, nil
},
}
svc := createTestServiceWithAgent(repo, newMockStepGetter(), newMockTaskRepo(), agentMgr)
svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
_, err = svc.PromptTask(ctx, "t1", "s1", "hello", "", false, nil, false)
if err == nil {
t.Fatal("expected the exhausted cold-resume retry to fail")
}
if errors.Is(err, ErrSessionRuntimeUnavailable) {
t.Fatalf("a genuine resume failure must not classify as ErrSessionRuntimeUnavailable "+
"(it would be silently queued with no visible error): %v", err)
}
}
// executorLookupErrorRepo wraps the real repo but forces
// GetExecutorRunningBySessionID to fail with a non-not-found error, standing
// in for a transient DB/context/deserialization failure.
type executorLookupErrorRepo struct {
*sqliterepo.Repository
err error
}
func (r *executorLookupErrorRepo) GetExecutorRunningBySessionID(context.Context, string) (*models.ExecutorRunning, error) {
return nil, r.err
}
// TestPromptTask_ExecutorLookupErrorNotClassifiedRuntimeUnavailable pins a
// second negative case alongside TestPromptTask_ExhaustedColdResumeNotClassifiedRuntimeUnavailable:
// a genuine executor-lookup failure (DB error, context cancellation, metadata
// deserialization) is not the "no row yet" shape errSessionAwaitingRuntimeLaunch
// exists for. It must stay visible instead of being classified as
// ErrSessionRuntimeUnavailable and silently queued with no future
// agent.boot_ready to drain it.
func TestPromptTask_ExecutorLookupErrorNotClassifiedRuntimeUnavailable(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput)
svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
svc.executor = executor.NewExecutor(&mockAgentManager{repoForExecutionLookup: repo}, repo, testLogger(), executor.ExecutorConfig{})
svc.repo = &executorLookupErrorRepo{Repository: repo, err: errors.New("db unavailable")}
_, err := svc.PromptTask(ctx, "t1", "s1", "hello", "", false, nil, false)
if err == nil {
t.Fatal("expected a visible error for a genuine executor lookup failure")
}
if errors.Is(err, ErrSessionRuntimeUnavailable) {
t.Fatalf("a genuine executor lookup error must not classify as ErrSessionRuntimeUnavailable "+
"(it would be silently queued with no future drain trigger): %v", err)
}
}
@@ -1,6 +1,7 @@
package orchestrator
import (
"context"
"errors"
"sync/atomic"
@@ -294,6 +295,22 @@ func (s *Service) clearQueuedDispatchInFlightIfCurrent(
}
}
func (s *Service) markQueuedDispatchDrainPending(sessionID string) {
if sessionID != "" {
s.queuedDispatchDrainPending.Store(sessionID, struct{}{})
}
}
func (s *Service) drainQueuedDispatchIfPending(sessionID string) {
if sessionID == "" {
return
}
if _, pending := s.queuedDispatchDrainPending.LoadAndDelete(sessionID); !pending {
return
}
s.drainQueuedMessageForPromptableSession(context.Background(), sessionID)
}
// releaseQueuedDispatchPendingIfCurrent is used by the fast prompt-claim
// helpers. Ownership has already moved to the accepted map, so they must not
// clear the accepted marker while the agent turn is still running.
+51 -12
View File
@@ -979,6 +979,10 @@ type Service struct {
// arbitrated through cancelInFlight.
dispatchingQueued sync.Map
acceptedQueuedDispatch sync.Map
// queuedDispatchDrainPending records a boot-ready event that arrived while
// another queued dispatch still owned the session. The worker consumes the
// marker after it clears that reservation and starts one deferred drain.
queuedDispatchDrainPending sync.Map // map[sessionID]struct{}
// afterReadyLifecycleReservation is a deterministic test seam for the
// narrow interval after handleAgentReady releases its per-session guard
@@ -3413,6 +3417,29 @@ func (s *Service) GetMessageQueue() *messagequeue.Service {
return s.messageQueue
}
type sessionPromptAdmissionContextKey struct{}
func (s *Service) withSessionPromptAdmission(
ctx context.Context,
sessionID string,
fn func(context.Context) error,
) error {
if fn == nil {
return errors.New("session prompt admission callback is nil")
}
if sessionID != "" {
if heldSessionID, _ := ctx.Value(sessionPromptAdmissionContextKey{}).(string); heldSessionID == sessionID {
return fn(ctx)
}
}
if s.messageQueue == nil {
return fn(ctx)
}
return s.messageQueue.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
return fn(context.WithValue(admittedCtx, sessionPromptAdmissionContextKey{}, sessionID))
})
}
// QueueUserPrompt persists a prompt that must wait for workflow WIP admission.
// The user message row is already written by the WebSocket handler, so the
// queue marker prevents the drain path from creating a duplicate row.
@@ -3434,18 +3461,30 @@ func (s *Service) QueueUserPrompt(
if userMessageRecorded {
queueMetadata[metaKeyUserMessageRecorded] = true
}
if _, err := s.messageQueue.QueueMessageWithMetadata(
ctx,
sessionID,
taskID,
prompt,
model,
messagequeue.QueuedByUser,
planMode,
toQueuedAttachments(attachments),
queueMetadata,
); err != nil {
return fmt.Errorf("queue user prompt: %w", err)
if err := s.withSessionPromptAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
session, err := s.repo.GetTaskSession(admittedCtx, sessionID)
if err != nil {
return fmt.Errorf("load session before queueing user prompt: %w", err)
}
if session == nil || isTerminalSessionState(session.State) {
return ErrSessionNotPromptable
}
if _, err := s.messageQueue.QueueMessageWithMetadata(
admittedCtx,
sessionID,
taskID,
prompt,
model,
messagequeue.QueuedByUser,
planMode,
toQueuedAttachments(attachments),
queueMetadata,
); err != nil {
return fmt.Errorf("queue user prompt: %w", err)
}
return nil
}); err != nil {
return err
}
s.publishQueueStatusEvent(ctx, sessionID)
@@ -79,6 +79,28 @@ var ErrAgentPromptInProgress = errors.New("agent is currently processing a promp
var ErrAgentNotReadyForPrompt = errors.New("agent not ready for prompt")
var ErrSessionResetInProgress = errors.New("session reset in progress")
// ErrSessionRuntimeUnavailable is returned by promptTask when
// ensureSessionRunning fails for the single reason that is safe to treat as
// "not yet launched" (see errSessionAwaitingRuntimeLaunch) — e.g. a session
// promoted by a workflow step move whose runtime has not finished launching
// yet. Callers can safely queue the prompt for delivery once the runtime
// comes up instead of reporting it as failed, because nothing was
// dispatched. Every other ensureSessionRunning failure (a real launch
// error, an office-scheduler refusal, an exhausted resume retry) keeps
// surfacing as a plain, visible error instead.
var ErrSessionRuntimeUnavailable = errors.New("session runtime unavailable")
// errSessionAwaitingRuntimeLaunch marks attemptColdResume's "prepared but
// not launched yet" outcome: the session has no executors_running row
// because the launch that creates one has not run yet, not because a launch
// attempt failed. This is the only ensureSessionRunning failure narrow
// enough to be classified ErrSessionRuntimeUnavailable — a session that
// tried and failed to come up (bad SSH handshake, exhausted cold-resume
// retries, an office-scheduler refusal) must not be silently queued with no
// visible error and no log, since there is no future launch that will drain
// the queue for it.
var errSessionAwaitingRuntimeLaunch = errors.New("session is not resumable: no executor record")
type primarySessionTaskStateUpdater interface {
UpdateTaskStateIfPrimarySessionState(
ctx context.Context,
@@ -2631,8 +2653,11 @@ func (s *Service) attemptColdResume(
}
running, lookupErr := s.repo.GetExecutorRunningBySessionID(ctx, sessionID)
if lookupErr != nil || running == nil {
return false, fmt.Errorf("session is not resumable: no executor record (state: %s)", session.State)
if lookupErr != nil && !errors.Is(lookupErr, models.ErrExecutorRunningNotFound) {
return false, fmt.Errorf("get executor running row: %w", lookupErr)
}
if running == nil {
return false, fmt.Errorf("%w (state: %s)", errSessionAwaitingRuntimeLaunch, session.State)
}
s.noteMissingWorktreesBeforeResume(sessionID, session)
@@ -4325,6 +4350,9 @@ func (s *Service) promptTask(ctx context.Context, taskID, sessionID string, prom
resumedForPrompt := !hadExecutionBeforeEnsure
if err := s.ensureSessionRunning(ctx, sessionID, session); err != nil {
s.releaseForegroundClaimOnFailure(ctx, taskID, sessionID, foregroundClaim)
if errors.Is(err, errSessionAwaitingRuntimeLaunch) {
return nil, fmt.Errorf("%w: failed to ensure session is running: %w", ErrSessionRuntimeUnavailable, err)
}
return nil, fmt.Errorf("failed to ensure session is running: %w", err)
}
@@ -972,12 +972,66 @@ func (h *MessageHandlers) forwardMessageAsPrompt(
// Don't create a prompt error message if the agent itself reported the error.
// The agent failure path (handleAgentFailed) already sets the session to FAILED
// with the error_message, which the UI displays via agent-status.
if !isAgentReportedError(err) {
if !isAgentReportedError(err) &&
!h.queuePromptIfRuntimeUnavailable(ctx, taskID, sessionID, content, model, planMode, attachments, err) {
h.createPromptErrorMessage(ctx, taskID, sessionID, err)
}
}
}
// queuePromptIfRuntimeUnavailable handles the window where a workflow step
// move has promoted a new primary session but its runtime has not finished
// launching: PromptTask fails with orchestrator.ErrSessionRuntimeUnavailable
// before anything reached the agent, so the message is safe to queue for
// delivery once the runtime comes up instead of being reported as failed.
// Returns true when the message was queued, so the caller must not also
// report promptErr as an error.
func (h *MessageHandlers) queuePromptIfRuntimeUnavailable(
ctx context.Context,
taskID, sessionID, content, model string,
planMode bool,
attachments []v1.MessageAttachment,
promptErr error,
) bool {
if !errors.Is(promptErr, orchestrator.ErrSessionRuntimeUnavailable) {
return false
}
session, err := h.service.GetTaskSession(ctx, sessionID)
if err != nil || session == nil || isTerminalSessionState(session.State) {
return false
}
// wsAddMessage already ran ProcessOnTurnStart synchronously for this prompt
// before the runtime-unavailable failure was even known; tag the queued
// entry so the drain path (executeQueuedMessageWithReservation) does not
// fire on_turn_start a second time on the replacement session.
queueMetadata := map[string]interface{}{orchestrator.MetaKeyTurnStartAlreadyProcessed: true}
if queueErr := h.orchestrator.QueueUserPrompt(
ctx, taskID, sessionID, content, model, planMode, attachments, queueMetadata, true,
); queueErr != nil {
h.logger.Warn("failed to queue prompt after runtime-unavailable prompt failure",
zap.String("task_id", taskID),
zap.String("session_id", sessionID),
zap.Error(queueErr))
return false
}
h.logger.Warn("queued prompt for delivery once session runtime finishes launching",
zap.String("task_id", taskID),
zap.String("session_id", sessionID),
zap.Error(promptErr))
return true
}
// isTerminalSessionState reports whether a session in this state can no
// longer accept a queued prompt for later delivery.
func isTerminalSessionState(state models.TaskSessionState) bool {
switch state {
case models.TaskSessionStateFailed, models.TaskSessionStateCancelled, models.TaskSessionStateCompleted:
return true
default:
return false
}
}
// isAgentReportedError returns true when the error originated from the agent's
// own error event (surfaced via waitForPromptDone with the ErrAgentReported
// sentinel wrapped in).
@@ -0,0 +1,107 @@
package handlers
import (
"context"
"errors"
"fmt"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/kandev/kandev/internal/orchestrator"
"github.com/kandev/kandev/internal/task/models"
)
// TestForwardMessageAsPrompt_QueuesMessageDroppedDuringProfileSwitchWindow is
// the regression test for the message-drop card: a message sent while a
// workflow step move is promoting a new primary session to a different agent
// profile used to be silently discarded. PromptTask fails with exactly the
// production shape — ensureSessionRunning's "no executor record" error
// wrapped in orchestrator.ErrSessionRuntimeUnavailable, matching neither of
// handlePromptWithResume's retry sentinels — so the message must be queued
// against the session instead of reported as failed.
func TestForwardMessageAsPrompt_QueuesMessageDroppedDuringProfileSwitchWindow(t *testing.T) {
notResumableErr := errors.New("session is not resumable: no executor record (state: WAITING_FOR_INPUT)")
promptErr := fmt.Errorf(
"%w: failed to ensure session is running: %w", orchestrator.ErrSessionRuntimeUnavailable, notResumableErr,
)
repo := &resumeRetryRepo{
sessionStateSequencer: sessionStateSequencer{
states: []models.TaskSessionState{models.TaskSessionStateWaitingForInput},
},
}
orch := &resumeRetryOrchestrator{promptErr: promptErr}
h := newTestMessageHandlersWithOrchestrator(t, repo, orch)
h.forwardMessageAsPrompt(
context.Background(), "task-1", "session-1", "profile-1", "continue",
"", false, nil, nil, false, "",
)
assert.Equal(t, 1, orch.promptCalls, "PromptTask must not be retried: neither retry sentinel matches this error class")
assert.Equal(t, 0, orch.resumeCalls)
assert.Equal(t, 1, orch.queuePromptCalls, "the message must be queued for delivery once the runtime comes up")
assert.Equal(t, true, orch.queueMetadata[orchestrator.MetaKeyTurnStartAlreadyProcessed], "queued retry must preserve the completed turn-start admission")
assert.Empty(t, repo.createdMessages, "a queued message must not also surface as a dropped-message error")
}
// TestForwardMessageAsPrompt_RuntimeUnavailableTerminalSessionSurfacesError
// ensures the queue fallback does not apply once the session has reached a
// terminal state: there will be no future runtime boot to drain the queue
// against, so the message must surface as an error like today.
func TestForwardMessageAsPrompt_RuntimeUnavailableTerminalSessionSurfacesError(t *testing.T) {
notResumableErr := errors.New("session is not resumable: no executor record (state: FAILED)")
promptErr := fmt.Errorf(
"%w: failed to ensure session is running: %w", orchestrator.ErrSessionRuntimeUnavailable, notResumableErr,
)
repo := &resumeRetryRepo{
sessionStateSequencer: sessionStateSequencer{
states: []models.TaskSessionState{models.TaskSessionStateFailed},
},
}
orch := &resumeRetryOrchestrator{promptErr: promptErr}
h := newTestMessageHandlersWithOrchestrator(t, repo, orch)
h.forwardMessageAsPrompt(
context.Background(), "task-1", "session-1", "profile-1", "continue",
"", false, nil, nil, false, "",
)
assert.Equal(t, 0, orch.queuePromptCalls, "a terminal session must not be queued")
require.Len(t, repo.createdMessages, 1, "a terminal session must still surface the failure")
assert.Contains(t, repo.createdMessages[0].Content, "Failed to send message to agent")
}
// TestForwardMessageAsPrompt_RuntimeUnavailableQueueFailureSurfacesError
// covers QueueUserPrompt itself failing (e.g. the orchestrator has no
// message queue configured): the original prompt error must still surface,
// since the message was neither delivered nor queued.
func TestForwardMessageAsPrompt_RuntimeUnavailableQueueFailureSurfacesError(t *testing.T) {
notResumableErr := errors.New("session is not resumable: no executor record (state: WAITING_FOR_INPUT)")
promptErr := fmt.Errorf(
"%w: failed to ensure session is running: %w", orchestrator.ErrSessionRuntimeUnavailable, notResumableErr,
)
repo := &resumeRetryRepo{
sessionStateSequencer: sessionStateSequencer{
states: []models.TaskSessionState{models.TaskSessionStateWaitingForInput},
},
}
orch := &resumeRetryOrchestrator{
promptErr: promptErr,
queuePromptErr: errors.New("message queue is not configured"),
}
h := newTestMessageHandlersWithOrchestrator(t, repo, orch)
h.forwardMessageAsPrompt(
context.Background(), "task-1", "session-1", "profile-1", "continue",
"", false, nil, nil, false, "",
)
assert.Equal(t, 1, orch.queuePromptCalls, "the queue attempt must still be made")
require.Len(t, repo.createdMessages, 1, "a failed queue attempt must fall back to surfacing the error")
assert.Contains(t, repo.createdMessages[0].Content, "Failed to send message to agent")
}
@@ -37,12 +37,15 @@ func (r *resumeRetryRepo) CreateMessage(_ context.Context, message *models.Messa
// additionally fails the second (retry) call, so tests can assert that a
// failed retry still surfaces the original error rather than the retry's own.
type resumeRetryOrchestrator struct {
promptErr error
retryPromptErr error
promptCalls int
resumeCalls int
resumeErr error
callOrder []string
promptErr error
retryPromptErr error
promptCalls int
resumeCalls int
resumeErr error
callOrder []string
queuePromptCalls int
queuePromptErr error
queueMetadata map[string]interface{}
}
func (o *resumeRetryOrchestrator) PromptTask(
@@ -75,8 +78,10 @@ func (o *resumeRetryOrchestrator) ProcessOnTurnStart(context.Context, string, st
return orchestrator.ProcessOnTurnStartResult{}, nil
}
func (*resumeRetryOrchestrator) QueueUserPrompt(context.Context, string, string, string, string, bool, []v1.MessageAttachment, map[string]interface{}, bool) error {
return nil
func (o *resumeRetryOrchestrator) QueueUserPrompt(_ context.Context, _, _, _, _ string, _ bool, _ []v1.MessageAttachment, metadata map[string]interface{}, _ bool) error {
o.queuePromptCalls++
o.queueMetadata = metadata
return o.queuePromptErr
}
func (o *resumeRetryOrchestrator) StepRequiresCompletionSignal(context.Context, string) bool {
@@ -160,14 +165,18 @@ func TestForwardMessageAsPrompt_RetriesOnceWhenAgentNotReadyAfterResume(t *testi
assert.Empty(t, repo.createdMessages, "a successful automatic retry must not surface a 'Request timed out' error message to the user")
}
// TestForwardMessageAsPrompt_SurfacesErrorWhenResumeRetryAlsoFails ensures
// the automatic retry does not swallow a genuine, non-recoverable failure:
// if ResumeTaskSession itself fails, the original readiness error must still
// reach the user via createPromptErrorMessage.
func TestForwardMessageAsPrompt_SurfacesErrorWhenResumeRetryAlsoFails(t *testing.T) {
// TestForwardMessageAsPrompt_QueuesInsteadOfSurfacingWhenResumeRetryAlsoFails
// covers a session-runtime-unavailable failure (orchestrator.
// ErrSessionRuntimeUnavailable, the class ensureSessionRunning returns when a
// workflow step move promoted a new primary session whose runtime has not
// finished launching) where the resume attempt itself also fails. Nothing
// ever reached the agent, so the message is queued for delivery once the
// runtime comes up rather than reported as failed — see
// queuePromptIfRuntimeUnavailable.
func TestForwardMessageAsPrompt_QueuesInsteadOfSurfacingWhenResumeRetryAlsoFails(t *testing.T) {
readinessErr := fmt.Errorf("%w: %w", orchestrator.ErrAgentNotReadyForPrompt, context.DeadlineExceeded)
resumeErr := fmt.Errorf("agent not ready after resume: %w", readinessErr)
promptErr := fmt.Errorf("failed to ensure session is running: %w", resumeErr)
promptErr := fmt.Errorf("%w: failed to ensure session is running: %w", orchestrator.ErrSessionRuntimeUnavailable, resumeErr)
repo := &resumeRetryRepo{
sessionStateSequencer: sessionStateSequencer{
@@ -187,6 +196,39 @@ func TestForwardMessageAsPrompt_SurfacesErrorWhenResumeRetryAlsoFails(t *testing
assert.Equal(t, 1, orch.promptCalls, "PromptTask must not be retried when the resume itself fails")
assert.Equal(t, 1, orch.resumeCalls)
assert.Equal(t, 1, orch.queuePromptCalls, "a pre-dispatch runtime-unavailable failure must be queued, not dropped")
assert.Empty(t, repo.createdMessages, "a queued message must not also surface as an error")
}
// TestForwardMessageAsPrompt_SurfacesErrorWhenResumeRetryFailsForNonRuntimeError
// is the sibling of the queue-instead-of-drop case above: a genuinely
// unrecoverable failure that is NOT classed as ErrSessionRuntimeUnavailable
// (a plain readiness timeout, not the pre-dispatch launch-window failure)
// must still reach the user via createPromptErrorMessage when the resume
// attempt also fails.
func TestForwardMessageAsPrompt_SurfacesErrorWhenResumeRetryFailsForNonRuntimeError(t *testing.T) {
readinessErr := fmt.Errorf("%w: %w", orchestrator.ErrAgentNotReadyForPrompt, context.DeadlineExceeded)
promptErr := fmt.Errorf("agent not ready after resume: %w", readinessErr)
repo := &resumeRetryRepo{
sessionStateSequencer: sessionStateSequencer{
states: []models.TaskSessionState{models.TaskSessionStateWaitingForInput},
},
}
orch := &resumeRetryOrchestrator{
promptErr: promptErr,
resumeErr: fmt.Errorf("resume: no executor record"),
}
h := newTestMessageHandlersWithOrchestrator(t, repo, orch)
h.forwardMessageAsPrompt(
context.Background(), "task-1", "session-1", "profile-1", "continue",
"", false, nil, nil, false, "",
)
assert.Equal(t, 1, orch.promptCalls, "PromptTask must not be retried when the resume itself fails")
assert.Equal(t, 1, orch.resumeCalls)
assert.Equal(t, 0, orch.queuePromptCalls, "a non-runtime-unavailable failure must not be queued")
require.Len(t, repo.createdMessages, 1, "a genuinely unrecoverable failure must still surface an error message")
assert.Contains(t, repo.createdMessages[0].Content, "Request timed out")
}