diff --git a/apps/backend/internal/backendapp/adapters_plugin_messenger.go b/apps/backend/internal/backendapp/adapters_plugin_messenger.go index 81687040e0..a2839f8dcd 100644 --- a/apps/backend/internal/backendapp/adapters_plugin_messenger.go +++ b/apps/backend/internal/backendapp/adapters_plugin_messenger.go @@ -33,6 +33,7 @@ type messengerTaskSvc interface { type messengerOrch interface { GetMessageQueue() *messagequeue.Service + PublishQueueStatusEvent(ctx context.Context, sessionID string) StartCreatedSession(ctx context.Context, taskID, sessionID, agentProfileID, prompt string, skipMessageRecord, planMode, autoStart bool, attachments []v1.MessageAttachment, references []v1.EntityReference) (*orchexecutor.TaskExecution, error) PromptTask(ctx context.Context, taskID, sessionID, prompt, model string, planMode bool, attachments []v1.MessageAttachment, dispatchOnly bool) (*orchestrator.PromptResult, error) ResumeTaskSession(ctx context.Context, taskID, sessionID string) (*orchexecutor.TaskExecution, error) @@ -96,7 +97,6 @@ func (a pluginsTaskMessengerAdapter) resolveSession(ctx context.Context, taskID, } // queueMessage appends the prompt to a running/starting session's FIFO queue -// for delivery at its next turn boundary. func (a pluginsTaskMessengerAdapter) queueMessage(ctx context.Context, taskID string, session *taskmodels.TaskSession, text string, metadata map[string]interface{}) (plugins.PluginMessageResult, error) { queue := a.orch.GetMessageQueue() if queue == nil { @@ -105,6 +105,7 @@ func (a pluginsTaskMessengerAdapter) queueMessage(ctx context.Context, taskID st if _, err := queue.QueueMessageWithMetadata(ctx, session.ID, taskID, text, "", messagequeue.QueuedByUser, false, nil, metadata); err != nil { return plugins.PluginMessageResult{}, fmt.Errorf("failed to queue message: %w", err) } + a.orch.PublishQueueStatusEvent(context.WithoutCancel(ctx), session.ID) return plugins.PluginMessageResult{SessionID: session.ID, Status: "queued"}, nil } diff --git a/apps/backend/internal/backendapp/adapters_plugin_messenger_test.go b/apps/backend/internal/backendapp/adapters_plugin_messenger_test.go index ec160a52bc..2d4a7e37d3 100644 --- a/apps/backend/internal/backendapp/adapters_plugin_messenger_test.go +++ b/apps/backend/internal/backendapp/adapters_plugin_messenger_test.go @@ -60,12 +60,19 @@ type fakeMessengerOrch struct { startCalls int promptCalls int resumeCalls int + queueStatusCalls int + queueStatusCtx context.Context promptErr error promptFailFirstOnly bool } func (f *fakeMessengerOrch) GetMessageQueue() *messagequeue.Service { return f.queue } +func (f *fakeMessengerOrch) PublishQueueStatusEvent(ctx context.Context, _ string) { + f.queueStatusCalls++ + f.queueStatusCtx = ctx +} + func (f *fakeMessengerOrch) StartCreatedSession(_ context.Context, _, _, _, _ string, _, _, _ bool, _ []v1.MessageAttachment, _ []v1.EntityReference) (*orchexecutor.TaskExecution, error) { f.startCalls++ return &orchexecutor.TaskExecution{}, nil @@ -107,6 +114,20 @@ func TestPluginsMessenger_RunningSessionQueues(t *testing.T) { require.Equal(t, 1, orch.queue.GetStatus(context.Background(), "s1").Count, "message should be enqueued") require.Nil(t, tasks.created, "queued path records via the queue, not CreateMessage") require.Zero(t, orch.startCalls+orch.promptCalls) + require.Equal(t, 1, orch.queueStatusCalls, "queued message should publish queue status") +} + +func TestPluginsMessenger_QueueStatusSurvivesCancelledRequest(t *testing.T) { + tasks := &fakeMessengerTaskSvc{primary: &taskmodels.TaskSession{ID: "s1", TaskID: "t1", State: taskmodels.TaskSessionStateRunning}} + orch := &fakeMessengerOrch{} + a := newMessengerAdapter(t, tasks, orch) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := a.SendMessage(ctx, "t1", "", "do the thing", "plugin:p") + require.NoError(t, err) + require.NotNil(t, orch.queueStatusCtx) + require.NoError(t, orch.queueStatusCtx.Err(), "queue status publication must not inherit a cancelled request") } func TestPluginsMessenger_CreatedSessionStarts(t *testing.T) { diff --git a/apps/backend/internal/backendapp/gateway.go b/apps/backend/internal/backendapp/gateway.go index 832cfac192..b0558b8e9a 100644 --- a/apps/backend/internal/backendapp/gateway.go +++ b/apps/backend/internal/backendapp/gateway.go @@ -120,6 +120,7 @@ func provideGateway( referenceValidator entityrefs.SubmissionValidator, authSvc *auth.Service, dataDir string, + registerCleanup func(func() error), lspMaxConnections ...int, ) (*gateways.Gateway, *notificationservice.Service, *notificationcontroller.Controller, *terminalservice.Service, error) { gateway, err := gateways.Provide(log) @@ -165,7 +166,19 @@ func provideGateway( referenceValidator, ) queueHandlers.SetAttachmentClaimer(taskSvc) + queueHandlers.Start(ctx) + if registerCleanup != nil { + registerCleanup(func() error { + queueHandlers.Stop() + return nil + }) + } queueHandlers.RegisterHandlers(gateway.Dispatcher) + if queue := orchestratorSvc.GetMessageQueue(); queue != nil { + gateway.Hub.SetClientDisconnectListener(func(connectionID string) { + queue.ReleaseEditLeasesForConnection(connectionID) + }) + } if lifecycleMgr != nil && agentRegistry != nil { agentCtrl := agentcontroller.NewController(lifecycleMgr, agentRegistry) diff --git a/apps/backend/internal/backendapp/main.go b/apps/backend/internal/backendapp/main.go index 197d4428ed..6bfd670629 100644 --- a/apps/backend/internal/backendapp/main.go +++ b/apps/backend/internal/backendapp/main.go @@ -943,6 +943,7 @@ func startGatewayAndServe( // be resolved while authentication is enforced. services.Auth, cfg.ResolvedHomeDir(), + func(fn func() error) { addCleanup(fn) }, cfg.Limits.LSPMaxConnections, ) if terminalSvc != nil { diff --git a/apps/backend/internal/backendapp/orchestrator.go b/apps/backend/internal/backendapp/orchestrator.go index 833d35da51..0462bc694d 100644 --- a/apps/backend/internal/backendapp/orchestrator.go +++ b/apps/backend/internal/backendapp/orchestrator.go @@ -142,6 +142,10 @@ func provideOrchestrator( } orchestratorSvc.SetAttachmentReader(taskSvc.AttachmentService()) orchestratorSvc.SetLaunchAttachmentClaimer(taskSvc) + if attachmentSvc := taskSvc.AttachmentService(); attachmentSvc != nil { + orchestratorSvc.SetSessionAttachmentCleaner(taskSvc) + } + orchestratorSvc.SetSessionAttachmentTransferer(taskSvc) orchestratorSvc.SetTitleBranchRuntime(lifecycleMgr) if githubSvc != nil { orchestratorSvc.SetTaskGitCredentialPolicyResolver(githubExecutorCredentialPolicyAdapter{service: githubSvc}) diff --git a/apps/backend/internal/gateway/websocket/client.go b/apps/backend/internal/gateway/websocket/client.go index 21291c6ecc..3cf1189cf6 100644 --- a/apps/backend/internal/gateway/websocket/client.go +++ b/apps/backend/internal/gateway/websocket/client.go @@ -101,10 +101,11 @@ func NewClient(id string, identity authn.Identity, conn *websocket.Conn, hub *Hu } // dispatchContext returns the hub's lifetime context carrying this client's -// identity, so dispatched RPC handlers (workspace.list, task CRUD, …) and -// subscription checks apply the same per-user scoping as HTTP requests. +// identity and server-assigned connection ID. The latter is used by +// connection-bound queue edit leases and never comes from the request payload. func (c *Client) dispatchContext() context.Context { ctx := c.hub.DispatchContext() + ctx = ws.WithConnectionID(ctx, c.ID) if c.identity.UserID != "" { ctx = authn.WithIdentity(ctx, c.identity) } diff --git a/apps/backend/internal/gateway/websocket/hub.go b/apps/backend/internal/gateway/websocket/hub.go index 53ce1d3583..2d78145033 100644 --- a/apps/backend/internal/gateway/websocket/hub.go +++ b/apps/backend/internal/gateway/websocket/hub.go @@ -54,6 +54,11 @@ type Hub struct { sessionGitDataProvider SessionGitDataProvider userSubscriptionListeners []func(userID string) + // clientDisconnectListener releases connection-bound resources after a + // client is removed from the hub. It runs asynchronously so durable cleanup + // cannot block the hub event loop. + clientDisconnectListener func(connectionID string) + // sessionMode tracks per-session focus state and fires listeners when // effective mode (paused/slow/fast) transitions. See hub_session_mode.go. sessionMode *sessionModeTracker @@ -137,7 +142,9 @@ func (h *Hub) Run(ctx context.Context) { func (h *Hub) closeAllClients() { h.mu.Lock() metricClientIDs := make([]string, 0, len(h.systemMetricsSubscribers)) + disconnectedIDs := make([]string, 0, len(h.clients)) for client := range h.clients { + disconnectedIDs = append(disconnectedIDs, client.ID) if client.systemMetricsSubscribed { metricClientIDs = append(metricClientIDs, client.ID) client.systemMetricsSubscribed = false @@ -146,6 +153,7 @@ func (h *Hub) closeAllClients() { delete(h.clients, client) } tracker := h.metricsInterestTracker + listener := h.clientDisconnectListener h.taskSubscribers = make(map[string]map[*Client]bool) h.sessionSubscribers = make(map[string]map[*Client]bool) h.runSubscribers = make(map[string]map[*Client]bool) @@ -159,6 +167,12 @@ func (h *Hub) closeAllClients() { } } + if listener != nil { + for _, clientID := range disconnectedIDs { + go listener(clientID) + } + } + h.stopAllPendingTransitions() } @@ -205,12 +219,17 @@ func (h *Hub) removeClient(client *Client) { metricClientID = client.ID tracker = h.metricsInterestTracker } + listener := h.clientDisconnectListener h.mu.Unlock() if tracker != nil && metricClientID != "" { tracker.MetricsUnsubscribe(metricClientID) } + if listener != nil { + go listener(client.ID) + } + for _, sessionID := range dedupStrings(affectedSessions) { h.recomputeSessionMode(sessionID) } @@ -218,6 +237,14 @@ func (h *Hub) removeClient(client *Client) { h.logger.Debug("Client unregistered", zap.String("client_id", client.ID)) } +// SetClientDisconnectListener registers a callback for connection teardown. +// The callback runs once for each client that was present in the hub. +func (h *Hub) SetClientDisconnectListener(listener func(connectionID string)) { + h.mu.Lock() + defer h.mu.Unlock() + h.clientDisconnectListener = listener +} + func (h *Hub) SetSystemMetricsInterestTracker(tracker SystemMetricsInterestTracker) { h.mu.Lock() defer h.mu.Unlock() diff --git a/apps/backend/internal/gateway/websocket/hub_lifecycle_test.go b/apps/backend/internal/gateway/websocket/hub_lifecycle_test.go index 3fab84ae83..461d45e2b0 100644 --- a/apps/backend/internal/gateway/websocket/hub_lifecycle_test.go +++ b/apps/backend/internal/gateway/websocket/hub_lifecycle_test.go @@ -40,3 +40,30 @@ func TestHub_ClientLifecycleCallsReturnAfterShutdown(t *testing.T) { }) } } + +func TestHubClientDisconnectNotifiesListenerOnce(t *testing.T) { + h := newTestHub(t) + client := newTestClient("connection-a") + h.clients[client] = true + disconnected := make(chan string, 1) + h.SetClientDisconnectListener(func(connectionID string) { + disconnected <- connectionID + }) + + h.removeClient(client) + h.removeClient(client) + + select { + case connectionID := <-disconnected: + if connectionID != client.ID { + t.Fatalf("disconnect notification = %q, want %q", connectionID, client.ID) + } + case <-time.After(time.Second): + t.Fatal("disconnect listener was not called") + } + select { + case duplicate := <-disconnected: + t.Fatalf("duplicate disconnect notification = %q", duplicate) + case <-time.After(50 * time.Millisecond): + } +} diff --git a/apps/backend/internal/gateway/websocket/task_notifications.go b/apps/backend/internal/gateway/websocket/task_notifications.go index f4b8250ec9..4beb3ca492 100644 --- a/apps/backend/internal/gateway/websocket/task_notifications.go +++ b/apps/backend/internal/gateway/websocket/task_notifications.go @@ -11,11 +11,6 @@ import ( "go.uber.org/zap" ) -const ( - queueStatusScopeKey = "queue_status_scope" - queueStatusScopeTask = "task" -) - type TaskEventBroadcaster struct { hub *Hub subscriptions []bus.Subscription @@ -290,13 +285,14 @@ func (b *TaskEventBroadcaster) routeBroadcast( b.hub.BroadcastToWorkspaceOrDrop(workspaceID, msg) return nil case ws.ActionMessageQueueStatusChanged: - if sessionID != "" { - b.hub.BroadcastToSession(sessionID, msg) - return nil - } - if extractStringField(data, queueStatusScopeKey) == queueStatusScopeTask { + // Queue status is session-scoped. An unscoped event cannot be safely + // attributed after a session is deleted, so fail closed rather than + // falling through to the global broadcast path. + if sessionID == "" { return nil } + b.hub.BroadcastToSession(sessionID, msg) + return nil case ws.ActionExecutorPrepareProgress, ws.ActionExecutorPrepareCompleted: // Broadcast to the owning workspace's clients so prepare // progress/warnings are available when the user navigates to the diff --git a/apps/backend/internal/gateway/websocket/task_notifications_test.go b/apps/backend/internal/gateway/websocket/task_notifications_test.go index 7dde94c519..bc3388709b 100644 --- a/apps/backend/internal/gateway/websocket/task_notifications_test.go +++ b/apps/backend/internal/gateway/websocket/task_notifications_test.go @@ -291,6 +291,29 @@ func TestTaskEventBroadcaster_DropsUnscopedGitHubCIOptionsWhenAuthIsEnforced(t * default: } } +func TestTaskEventBroadcaster_DropsUnscopedQueueStatus(t *testing.T) { + hub := newTestHub(t) + hub.setAuthPolicy(AuthPolicy{Enforced: func() bool { return true }}) + msg, err := ws.NewNotification(ws.ActionMessageQueueStatusChanged, map[string]any{ + "task_id": "task-without-session", + }) + require.NoError(t, err) + broadcaster := &TaskEventBroadcaster{hub: hub, logger: testLogger()} + + require.NoError(t, broadcaster.routeBroadcast( + ws.ActionMessageQueueStatusChanged, + msg.Payload, + "", + "", + msg, + )) + + select { + case leaked := <-hub.broadcast: + t.Fatalf("unscoped queue status was globally broadcast: %s", leaked.Action) + default: + } +} func TestTaskEventBroadcaster_ScopesTypedGitHubTaskPRUpdateToOwningWorkspace(t *testing.T) { hub := newTestHub(t) diff --git a/apps/backend/internal/mcp/handlers/config_handlers_test.go b/apps/backend/internal/mcp/handlers/config_handlers_test.go index 915ee3a45b..aca38b793d 100644 --- a/apps/backend/internal/mcp/handlers/config_handlers_test.go +++ b/apps/backend/internal/mcp/handlers/config_handlers_test.go @@ -540,6 +540,67 @@ func TestDeferMoveTask_RollsBackHandoffWhenPendingMovePersistenceFails(t *testin assert.Equal(t, "preexisting", queue.calls[0].ID) } +type rollbackContextMessageQueuer struct { + recordingMessageQueuer + queued *messagequeue.QueuedMessage + rollbackCtx context.Context + removedEntry string + takeCalled bool +} + +func (r *rollbackContextMessageQueuer) QueueMessageWithMetadata( + ctx context.Context, + sessionID, taskID, content, model, userID string, + planMode bool, + attachments []messagequeue.MessageAttachment, + metadata map[string]interface{}, +) (*messagequeue.QueuedMessage, error) { + msg, err := r.recordingMessageQueuer.QueueMessageWithMetadata( + ctx, sessionID, taskID, content, model, userID, planMode, attachments, metadata, + ) + if err == nil { + msg.ID = "queued-handoff" + r.queued = msg + } + return msg, err +} + +func (r *rollbackContextMessageQueuer) TakeQueued(ctx context.Context, _ string) (*messagequeue.QueuedMessage, bool) { + r.rollbackCtx = ctx + r.takeCalled = true + if r.queued == nil { + return nil, false + } + msg := r.queued + r.queued = nil + return msg, true +} + +func (r *rollbackContextMessageQueuer) RemoveEntry(ctx context.Context, _ string, entryID string) error { + r.rollbackCtx = ctx + r.removedEntry = entryID + if r.queued != nil && r.queued.ID == entryID { + r.queued = nil + } + return nil +} + +func (r *rollbackContextMessageQueuer) RemoveEntryForSession( + ctx context.Context, + identity messagequeue.QueueSessionIdentity, + entryID string, +) (*messagequeue.QueueRemovalResult, error) { + r.rollbackCtx = ctx + r.removedEntry = entryID + if r.queued == nil || r.queued.ID != entryID || + r.queued.SessionID != identity.SessionID || r.queued.TaskID != identity.TaskID { + return nil, messagequeue.ErrEntryNotFound + } + removed := *r.queued + r.queued = nil + return &messagequeue.QueueRemovalResult{Removed: []messagequeue.QueuedMessage{removed}}, nil +} + // TestQueueMoveTaskPrompt_NilQueueReturnsError ensures the call is safe (no panic) // and surfaces a descriptive error so callers can fail fast instead of silently // dropping the user-supplied prompt. @@ -596,6 +657,41 @@ func TestQueueMoveTaskPrompt_QueuesWithExpectedFields(t *testing.T) { assert.Equal(t, "", got.Model) } +func TestApplyMoveTaskImmediateRollsBackQueueWithDetachedContext(t *testing.T) { + taskSvc, _ := newTestTaskService(t) + queue := &rollbackContextMessageQueuer{} + h := &Handlers{ + taskSvc: taskSvc, + messageQueue: queue, + logger: testLogger(t).WithFields(), + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + msg := makeWSMessage(t, ws.ActionMCPMoveTask, nil) + req := struct { + TaskID string `json:"task_id"` + WorkflowID string `json:"workflow_id"` + WorkflowStepID string `json:"workflow_step_id"` + Position int `json:"position"` + Prompt string `json:"prompt"` + SenderSessionID string `json:"sender_session_id"` + }{ + TaskID: "missing-task", + WorkflowID: "missing-workflow", + WorkflowStepID: "missing-step", + Prompt: "handoff", + } + session := &models.TaskSession{ID: "session-rollback", TaskID: req.TaskID} + + resp, err := h.applyMoveTaskImmediate(ctx, msg, req, session) + require.NoError(t, err) + assertWSError(t, resp, ws.ErrorCodeInternalError) + require.NotNil(t, queue.rollbackCtx) + assert.NoError(t, queue.rollbackCtx.Err()) + assert.Equal(t, "queued-handoff", queue.removedEntry) + assert.False(t, queue.takeCalled) + assert.Nil(t, queue.queued) +} func TestHandleDeleteTask_MissingTaskID(t *testing.T) { h := &Handlers{} msg := makeWSMessage(t, ws.ActionMCPDeleteTask, map[string]string{}) diff --git a/apps/backend/internal/mcp/handlers/config_task_handlers.go b/apps/backend/internal/mcp/handlers/config_task_handlers.go index 86a1ec6379..3818c88956 100644 --- a/apps/backend/internal/mcp/handlers/config_task_handlers.go +++ b/apps/backend/internal/mcp/handlers/config_task_handlers.go @@ -362,7 +362,6 @@ func (h *Handlers) queueMoveTaskPromptWithMoveID( if moveID != "" { metadata = map[string]interface{}{messagequeue.MetadataDeferredMoveID: moveID} } - var ( entry *messagequeue.QueuedMessage err error @@ -396,7 +395,8 @@ func (h *Handlers) rollbackMoveTaskPrompt(ctx context.Context, handoff *queuedMo if handoff == nil || h.messageQueue == nil { return } - if _, err := h.messageQueue.RemoveEntryForSession(ctx, handoff.identity, handoff.entryID); err != nil { + rollbackCtx := context.WithoutCancel(ctx) + if _, err := h.messageQueue.RemoveEntryForSession(rollbackCtx, handoff.identity, handoff.entryID); err != nil { h.logger.Error("move_task: failed to roll back queued hand-off prompt", zap.String("task_id", handoff.identity.TaskID), zap.String("session_id", handoff.identity.SessionID), @@ -405,7 +405,7 @@ func (h *Handlers) rollbackMoveTaskPrompt(ctx context.Context, handoff *queuedMo return } if queue, ok := h.messageQueue.(*messagequeue.Service); ok { - h.publishQueueStatusEvent(ctx, handoff.identity, queue) + h.publishQueueStatusEvent(rollbackCtx, handoff.identity, queue) } } diff --git a/apps/backend/internal/mcp/handlers/handlers.go b/apps/backend/internal/mcp/handlers/handlers.go index e0aa846b64..3fa335ee36 100644 --- a/apps/backend/internal/mcp/handlers/handlers.go +++ b/apps/backend/internal/mcp/handlers/handlers.go @@ -2988,6 +2988,7 @@ type taskMessageDispatchResult struct { } type taskMessageReviewRollback struct { + taskID string changed bool restoreTask bool taskState v1.TaskState @@ -3452,6 +3453,7 @@ func (h *Handlers) ensureTaskInProgressForTaskMessage(ctx context.Context, taskI return taskMessageReviewRollback{}, err } rollback := taskMessageReviewRollback{ + taskID: taskID, changed: true, restoreTask: true, taskState: task.State, @@ -3561,14 +3563,28 @@ func (h *Handlers) restoreSelectedTaskMessageSession(ctx context.Context, repo t return nil } selected, err := repo.GetTaskSession(ctx, rollback.selectedID) - if err != nil { + if err != nil && !errors.Is(err, models.ErrTaskSessionNotFound) { return err } if selected != nil && selected.State == models.TaskSessionStateCancelled { return errTaskMessageRollbackSuperseded } + taskID := rollback.taskID + if selected != nil && selected.TaskID != "" { + taskID = selected.TaskID + } + if taskID == "" { + if snapshot, ok := rollback.queues[rollback.selectedID]; ok { + for _, entry := range snapshot.entries { + if entry.TaskID != "" { + taskID = entry.TaskID + break + } + } + } + } if primaryID != "" && rollback.selectedID != primaryID { - if err := h.restoreTaskMessageQueueOwner(ctx, rollback.selectedID, primaryID); err != nil { + if err := h.restoreTaskMessageQueueOwner(ctx, taskID, rollback.selectedID, primaryID); err != nil { return err } } @@ -3610,10 +3626,12 @@ func (h *Handlers) restoreTaskMessageQueues(ctx context.Context, rollback taskMe if queue == nil { return nil } + restoreCtx := context.WithoutCancel(ctx) for sessionID, snapshot := range rollback.queues { - if err := h.restoreTaskMessageQueue(ctx, queue, sessionID, snapshot); err != nil { + if err := h.restoreTaskMessageQueue(restoreCtx, queue, sessionID, snapshot); err != nil { return err } + h.publishQueueStatusEvent(restoreCtx, snapshot.identity, queue) } return nil } @@ -3626,7 +3644,8 @@ func (h *Handlers) restoreTaskMessageQueue(ctx context.Context, queue *messagequ return queue.RestoreSessionForIdentity(ctx, snapshot.identity, cloneTaskMessageQueuedMessages(snapshot.entries), pendingMove) } -func (h *Handlers) restoreTaskMessageQueueOwner(ctx context.Context, selectedID, primaryID string) error { +//nolint:cyclop // Queue-owner restoration validates two identities and rolls back attachment transfer on failure. +func (h *Handlers) restoreTaskMessageQueueOwner(ctx context.Context, taskID, selectedID, primaryID string) error { queue := h.sessionLauncher.GetMessageQueue() if queue == nil { return nil @@ -3661,7 +3680,38 @@ func (h *Handlers) restoreTaskMessageQueueOwner(ctx context.Context, selectedID, primaryIdentity.SessionIncarnationID != primary.QueueIncarnationID { return messagequeue.ErrSessionIdentityMismatch } - return queue.TransferSessionIdentities(ctx, selectedIdentity, primaryIdentity) + transferCtx := context.WithoutCancel(ctx) + transferErr := queue.TransferSessionWithDurableAttachmentPreparation( + transferCtx, + taskID, + selectedID, + primaryID, + func(admittedCtx context.Context, attachmentIDs []string) error { + if h.taskSvc == nil { + return errors.New("session attachment transfer service is unavailable") + } + if err := h.taskSvc.TransferSessionMessageAttachments( + admittedCtx, taskID, selectedID, primaryID, attachmentIDs, + ); err != nil { + return fmt.Errorf("transfer session attachments: %w", err) + } + return nil + }, + func(rollbackCtx context.Context, attachmentIDs []string) error { + if h.taskSvc == nil { + return errors.New("session attachment transfer service is unavailable") + } + return h.taskSvc.TransferSessionMessageAttachments( + rollbackCtx, taskID, primaryID, selectedID, attachmentIDs, + ) + }, + ) + if transferErr != nil { + return transferErr + } + h.publishQueueStatusEvent(transferCtx, selectedIdentity, queue) + h.publishQueueStatusEvent(transferCtx, primaryIdentity, queue) + return nil } func restoreTaskMessageSessionSnapshot(ctx context.Context, repo taskMessageSessionRollbackRepository, rollback taskMessageSessionRollback) error { session, err := repo.GetTaskSession(ctx, rollback.sessionID) diff --git a/apps/backend/internal/mcp/handlers/message_task_queue_transfer_test.go b/apps/backend/internal/mcp/handlers/message_task_queue_transfer_test.go new file mode 100644 index 0000000000..e27154838a --- /dev/null +++ b/apps/backend/internal/mcp/handlers/message_task_queue_transfer_test.go @@ -0,0 +1,68 @@ +package handlers + +import ( + "context" + "strings" + "testing" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/task/models" + "github.com/kandev/kandev/internal/task/service" + "github.com/stretchr/testify/require" +) + +func TestRestoreTaskMessageQueueOwnerTransfersClaimedAttachments(t *testing.T) { + ctx := context.Background() + taskSvc, repo := newTestTaskService(t) + _, target, oldSession := seedTaskWithSession(t, taskSvc, repo, models.TaskSessionStateWaitingForInput) + newSession := taskSessionForQueueTransferTest(target.ID) + require.NoError(t, repo.CreateTaskSession(ctx, newSession)) + + attachments, err := service.NewAttachmentService(repo, t.TempDir(), nil, testLogger(t)) + require.NoError(t, err) + taskSvc.SetAttachmentService(attachments) + attachment, err := attachments.Stage(ctx, "user-a", "ws-1", "prompt.txt", "text/plain", "resource", "", strings.NewReader("prompt")) + require.NoError(t, err) + require.NoError(t, attachments.Claim(ctx, "user-a", "ws-1", target.ID, oldSession.ID, []string{attachment.ID})) + + h, orch := newMessageTaskHandler(t, taskSvc, repo) + _, err = orch.queue.QueueMessage(ctx, oldSession.ID, target.ID, "queued prompt", "", messagequeue.QueuedByUser, false, []messagequeue.MessageAttachment{{AttachmentID: attachment.ID}}) + require.NoError(t, err) + + require.NoError(t, h.restoreTaskMessageQueueOwner(ctx, target.ID, oldSession.ID, newSession.ID)) + + moved, err := repo.GetMessageAttachment(ctx, attachment.ID) + require.NoError(t, err) + require.Equal(t, newSession.ID, moved.SessionID) + status := orch.queue.GetStatus(ctx, newSession.ID) + require.Len(t, status.Entries, 1) + require.Equal(t, attachment.ID, status.Entries[0].Attachments[0].AttachmentID) +} + +func TestRestoreTaskMessageQueueOwnerFailsClosedWithoutAttachmentService(t *testing.T) { + ctx := context.Background() + taskSvc, repo := newTestTaskService(t) + _, target, oldSession := seedTaskWithSession(t, taskSvc, repo, models.TaskSessionStateWaitingForInput) + newSession := taskSessionForQueueTransferTest(target.ID) + require.NoError(t, repo.CreateTaskSession(ctx, newSession)) + + h, orch := newMessageTaskHandler(t, taskSvc, repo) + _, err := orch.queue.QueueMessage( + ctx, oldSession.ID, target.ID, "queued prompt", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "unavailable-attachment"}}, + ) + require.NoError(t, err) + + require.Error(t, h.restoreTaskMessageQueueOwner(ctx, target.ID, oldSession.ID, newSession.ID)) + status := orch.queue.GetStatus(ctx, oldSession.ID) + require.Len(t, status.Entries, 1) + require.Empty(t, orch.queue.GetStatus(ctx, newSession.ID).Entries) +} + +func taskSessionForQueueTransferTest(taskID string) *models.TaskSession { + return &models.TaskSession{ + ID: "replacement-session", + TaskID: taskID, + State: models.TaskSessionStateWaitingForInput, + } +} diff --git a/apps/backend/internal/mcp/handlers/queue_status_event_test.go b/apps/backend/internal/mcp/handlers/queue_status_event_test.go index 5e86aaa74c..d23c521aff 100644 --- a/apps/backend/internal/mcp/handlers/queue_status_event_test.go +++ b/apps/backend/internal/mcp/handlers/queue_status_event_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/task/models" "github.com/stretchr/testify/require" ) @@ -31,3 +32,45 @@ func TestPublishQueueStatusEventIncludesQueuePolicy(t *testing.T) { require.NotEqual(t, identity.SessionIncarnationID, data["status_epoch"]) require.Equal(t, true, data["auto_merge_available"]) } + +func TestPublishQueueStatusEventUsesOwningSessionTaskForMixedQueue(t *testing.T) { + ctx := context.Background() + taskSvc, repo := newTestTaskService(t) + seedMCPHandlerSession(t, repo, "task-owner", "session-mixed", models.TaskSessionStateWaitingForInput) + queue := messagequeue.NewServiceMemory(testLogger(t)) + eventBus := &mcpRecordingEventBus{} + handlers := &Handlers{taskSvc: taskSvc, eventBus: eventBus} + identity, err := queue.ResolveSessionIdentity(ctx, "task-owner", "session-mixed") + require.NoError(t, err) + + _, err = queue.QueueMessage(ctx, "session-mixed", "task-other", "from another task", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + _, err = queue.QueueMessage(ctx, "session-mixed", "task-owner", "from owner", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + + handlers.publishQueueStatusEvent(ctx, identity, queue) + + require.Len(t, eventBus.events, 1) + data, ok := eventBus.events[0].Data.(map[string]interface{}) + require.True(t, ok) + require.Equal(t, "task-owner", data["task_id"]) +} + +func TestQueueMoveTaskPromptPublishesQueueStatus(t *testing.T) { + ctx := context.Background() + queue := messagequeue.NewServiceMemory(testLogger(t)) + eventBus := &mcpRecordingEventBus{} + handlers := &Handlers{eventBus: eventBus, messageQueue: queue} + + identity, err := queue.ResolveSessionIdentity(ctx, "task-move", "session-move") + require.NoError(t, err) + _, err = handlers.queueMoveTaskPrompt(ctx, identity, "continue") + require.NoError(t, err) + + require.Len(t, eventBus.events, 1) + data, ok := eventBus.events[0].Data.(map[string]interface{}) + require.True(t, ok) + require.Equal(t, "session-move", data["session_id"]) + require.Equal(t, "task-move", data["task_id"]) + require.Equal(t, 1, data["count"]) +} diff --git a/apps/backend/internal/orchestrator/ci_automation_attempt.go b/apps/backend/internal/orchestrator/ci_automation_attempt.go index 2d91e8d3ca..b4b0c1cec6 100644 --- a/apps/backend/internal/orchestrator/ci_automation_attempt.go +++ b/apps/backend/internal/orchestrator/ci_automation_attempt.go @@ -103,10 +103,10 @@ func (s *Service) reconcileOrphanedCIAutoFixQueueEntries( continue } } - if ackErr := s.messageQueue.AcknowledgeQueued( + if removeErr := s.messageQueue.RemoveEntry( context.WithoutCancel(ctx), entry.SessionID, entry.ID, - ); ackErr != nil { - return fmt.Errorf("remove orphaned CI auto-fix queue entry %q: %w", entry.ID, ackErr) + ); removeErr != nil && !errors.Is(removeErr, messagequeue.ErrEntryNotFound) { + return fmt.Errorf("remove orphaned CI auto-fix queue entry %q: %w", entry.ID, removeErr) } } return nil diff --git a/apps/backend/internal/orchestrator/event_handlers_agent.go b/apps/backend/internal/orchestrator/event_handlers_agent.go index 6d109d38ad..5153998b21 100644 --- a/apps/backend/internal/orchestrator/event_handlers_agent.go +++ b/apps/backend/internal/orchestrator/event_handlers_agent.go @@ -184,6 +184,12 @@ func (s *Service) publishQueueStatusEvent(ctx context.Context, sessionID string) )) } +// PublishQueueStatusEvent publishes the authoritative queue snapshot after a +// queue mutation performed by an integration adapter. +func (s *Service) PublishQueueStatusEvent(ctx context.Context, sessionID string) { + s.publishQueueStatusEvent(context.WithoutCancel(ctx), sessionID) +} + func (s *Service) publishQueueStatusEventForIdentity( ctx context.Context, identity messagequeue.QueueSessionIdentity, @@ -221,19 +227,40 @@ func (s *Service) publishTaskQueueStatusEvent(ctx context.Context, taskID, sessi // Detach so a client disconnect cannot cancel projector recount (which would // leave queued_prompt_count stale on every live sidebar). ctx = context.WithoutCancel(ctx) + if sessionID != "" && s.messageQueue != nil { + _ = s.messageQueue.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + s.publishTaskQueueStatusEventSnapshot(admittedCtx, taskID, sessionID) + return nil + }) + return + } + s.publishTaskQueueStatusEventSnapshot(ctx, taskID, sessionID) +} + +func (s *Service) publishTaskQueueStatusEventSnapshot(ctx context.Context, taskID, sessionID string) { eventData := map[string]interface{}{ metaKeyTaskID: taskID, queueStatusScopeKey: queueStatusScopeTask, } if sessionID != "" && s.messageQueue != nil { - queueStatus, err := s.queueStatusSnapshot(ctx, sessionID) - if err != nil { - s.logger.Warn("snapshot task queue status event", - zap.String(metaKeyTaskID, taskID), - zap.String(metaKeySessionID, sessionID), - zap.Error(err)) - } else { + if s.repo == nil { + // Focused adapters can provide only the queue service. Use its + // authoritative in-memory snapshot when no task repository is + // available to resolve the session identity. + queueStatus := s.messageQueue.GetStatus(ctx, sessionID) + queueStatus.TaskID = taskID + queueStatus.SessionID = sessionID eventData = queueStatusEventData(queueStatus) + } else { + queueStatus, err := s.queueStatusSnapshot(ctx, sessionID) + if err != nil { + s.logger.Warn("snapshot task queue status event", + zap.String(metaKeyTaskID, taskID), + zap.String(metaKeySessionID, sessionID), + zap.Error(err)) + } else { + eventData = queueStatusEventData(queueStatus) + } } } s.logger.Debug("publishing task queue status changed event", @@ -254,9 +281,14 @@ func (s *Service) publishTaskQueueStatusEvent(ctx context.Context, taskID, sessi // FIFO priority; reserved lifecycle messages release their exact durable row // so a newer coalesced successor cannot overwrite the in-flight delivery. func (s *Service) requeueMessage(ctx context.Context, queuedMsg *messagequeue.QueuedMessage, queuedBy string) { - if queuedMsg.QueuedBy != "" { + coalesceKey := messageCoalesceKey(queuedMsg) + if queuedMsg.QueuedBy != "" && coalesceKey != "" { queuedBy = queuedMsg.QueuedBy } + if isLifecycleAutomationMessage(queuedMsg) && queuedMsg.IsReservedLifecycleDelivery() { + s.requeueLifecycleMessage(ctx, queuedMsg, queuedBy, coalesceKey) + return + } if err := s.messageQueue.RequeueAtHead(ctx, queuedMsg); err != nil { s.logger.Error("failed to requeue message at head", zap.String("session_id", queuedMsg.SessionID), @@ -621,10 +653,13 @@ func (s *Service) handleAgentReady(ctx context.Context, data watcher.AgentEventD zap.String("task_id", data.TaskID)) return } - // A passthrough lifecycle reservation cannot be executed while this ready - // handler holds the per-session guard: executeQueuedMessage performs its - // own final claim through that guard. Register the deferred dispatch before - // acquiring it, so LIFO defer ordering releases the guard first. + // A passthrough queue entry requiring the executor path (a lifecycle + // reservation or an attachment-bearing ordinary entry) cannot be executed + // while this ready handler holds the per-session guard: executeQueuedMessage + // performs its own final claim through that guard. Register the deferred + // dispatch before acquiring it, so LIFO defer ordering releases the guard + // first. Keeping execution in this handler preserves the legacy ready-event + // completion boundary while still using the normal delivery pipeline. var deferredLifecycleDispatch *messagequeue.QueuedMessage var deferredLifecycleReservation *queuedDispatchReservation defer func() { @@ -910,11 +945,28 @@ func (s *Service) handleAgentReady(ctx context.Context, data watcher.AgentEventD if err != nil || !autoRun || !exists { return } + // ReserveQueued changes the visible queue immediately. Publish before + // any passthrough delivery so ordinary, attachment, and lifecycle + // reservations all reach clients, including the deferred async path. + s.publishQueueStatusEvent(ctx, data.SessionID) + // A durable lifecycle reservation must take the same guarded delivery + // path as an ACP session. In particular, executeQueuedMessage performs + // the final active-task/session claim, records the visible message, and + // only acknowledges the reservation after Executor.Prompt accepts it. + // Executor.Prompt itself is passthrough-aware and writes to PTY stdin. + // Attachment-bearing ordinary entries use that same path so their + // attachments are materialized and included in the prompt. Ordinary + // text-only entries retain the historical direct PTY behavior below. if queuedMsg.IsDurableLifecycle() { deferredLifecycleReservation = s.markQueuedDispatchInFlightWithIdentityLocked(identity, queuedMsg.ID, queuedMsg) deferredLifecycleDispatch = queuedMsg return } + if len(queuedMsg.Attachments) > 0 { + deferredLifecycleReservation = s.markQueuedDispatchInFlightWithIdentityLocked(identity, queuedMsg.ID, queuedMsg) + deferredLifecycleDispatch = queuedMsg + return + } releaseReservation := func() { if releaseErr := s.messageQueue.ReleaseQueuedDeliveryForSession( context.WithoutCancel(ctx), identity, queuedMsg.ID, @@ -964,6 +1016,7 @@ func (s *Service) handleAgentReady(ctx context.Context, data watcher.AgentEventD releaseReservation() return } + return } if err := s.messageQueue.AcknowledgeQueuedForSession( context.WithoutCancel(ctx), identity, queuedMsg.ID, @@ -984,15 +1037,12 @@ func (s *Service) handleAgentReady(ctx context.Context, data watcher.AgentEventD const githubPRAutomationOrigin = "github_pr_automation" -// isLifecycleAutomationOrigin reports whether a queued message's "origin" -// metadata identifies it as a durable lifecycle prompt (GitHub PR or GitLab -// MR automation) rather than an ordinary queued message. Both producers -// share the same durable-queue contract (QueueLifecycleMessageWithCoalesceKey, -// AcknowledgeQueued on accepted delivery); this is the single place that -// recognizes the set of origins entitled to that treatment, so adding a -// future provider only needs a change here. -func isLifecycleAutomationOrigin(origin interface{}) bool { - return origin == githubPRAutomationOrigin || origin == mrAutomationOrigin || origin == ciAutomationOrigin +// isLifecycleAutomationMessage reports whether a queued message is a durable +// lifecycle prompt. The durable marker is authoritative for legacy rows whose +// provider origin is absent, while IsDurableLifecycle also recognizes known +// provider origins written by older versions. +func isLifecycleAutomationMessage(msg *messagequeue.QueuedMessage) bool { + return msg != nil && msg.IsDurableLifecycle() } func (s *Service) recordQueuedUserMessage(ctx context.Context, queuedMsg *messagequeue.QueuedMessage, attachments []v1.MessageAttachment) error { @@ -1045,7 +1095,7 @@ func (s *Service) executeQueuedMessageWithReservation( s.onQueuedMessageExecutionComplete() } }() - lifecyclePrompt := isLifecycleAutomationOrigin(queuedMsg.Metadata["origin"]) + lifecyclePrompt := isLifecycleAutomationMessage(queuedMsg) claimEntryID, handoffDone := s.claimQueuedMessageHandoff( promptCtx, callerSessionID, queuedMsg, reservation, @@ -1124,6 +1174,7 @@ func (s *Service) executeQueuedMessageWithReservation( afterClaim := s.queuedMessageAfterClaim( promptCtx, dispatchIdentity, queuedMsg, attachments, lifecyclePrompt, &userMessageRecorded, ) + afterDispatch := s.queuedMessageAfterDispatch(promptCtx, queuedMsg, lifecyclePrompt) _, err := s.promptTask(promptCtx, queuedMsg.TaskID, queuedMsg.SessionID, promptContent, queuedMsg.Model, queuedMsg.PlanMode, attachments, false, promptTaskOptions{ @@ -1133,6 +1184,7 @@ func (s *Service) executeQueuedMessageWithReservation( onAccepted: func(turnID string) { s.bindQueuedCIAutoFixAttempt(promptCtx, queuedMsg, turnID) }, + afterDispatch: afterDispatch, }) if err != nil { s.reconcileQueuedCIAutoFixDispatchFailure(promptCtx, queuedMsg) @@ -1239,9 +1291,10 @@ func (s *Service) finishQueuedMessageExecution( var reselected *lifecyclePromptReselectedError if errors.As(err, &reselected) { if !exactReservation { - queuedMsg.SessionID = reselected.sessionID + retry := *queuedMsg + retry.SessionID = reselected.sessionID if s.requeueLifecycleMessage( - ctx, queuedMsg, queuedMsg.QueuedBy, messageCoalesceKey(queuedMsg), + ctx, &retry, retry.QueuedBy, messageCoalesceKey(&retry), ) { s.acknowledgeLifecycleQueueEntry(ctx, reservedSessionID, queuedMsg) } @@ -1282,6 +1335,20 @@ func (s *Service) finishQueuedMessageExecution( } return } + var acceptedDispatch *acceptedPromptDispatchError + if errors.As(err, &acceptedDispatch) { + s.logger.Error("queued prompt was accepted but post-dispatch handling failed", + zap.String("session_id", callerSessionID), + zap.String("task_id", queuedMsg.TaskID), + zap.String("queue_id", queuedMsg.ID), + zap.Error(err)) + if lifecyclePrompt { + s.acknowledgeLifecycleQueueEntry(ctx, reservedSessionID, queuedMsg) + } else { + s.acknowledgeOrdinaryQueueEntry(ctx, reservedSessionID, queuedMsg) + } + return + } if err != nil { s.handleQueuedMessageExecutionError( ctx, callerSessionID, queuedMsg, reservation, lifecyclePrompt, userMessageRecorded, err, @@ -1294,6 +1361,8 @@ func (s *Service) finishQueuedMessageExecution( } else { s.acknowledgeLifecycleQueueEntry(ctx, reservedSessionID, queuedMsg) } + } else { + s.acknowledgeOrdinaryQueueEntry(ctx, reservedSessionID, queuedMsg) } } @@ -1312,9 +1381,11 @@ 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) || + passthroughAttachmentRecovery := !lifecyclePrompt && + len(queuedMsg.Attachments) > 0 && + s.agentManager != nil && + s.agentManager.IsPassthroughSession(ctx, queuedMsg.SessionID) + if passthroughAttachmentRecovery || lifecyclePrompt || errors.Is(err, errLifecyclePromptClaim) || errors.Is(err, errLifecyclePromptMessagePersistence) || isSessionBusyError(err) || isTransientPromptError(err) || manualRecovery || errors.Is(err, lifecycle.ErrCancelEscalated) || isSessionResetInProgressError(err) || @@ -1328,10 +1399,21 @@ func (s *Service) handleQueuedMessageExecutionError( zap.String("queue_id", queuedMsg.ID), zap.Bool("manual_recovery", manualRecovery)) switch { + case passthroughAttachmentRecovery && reservation != nil && reservation.identity.SessionIncarnationID != "": + s.restoreQueuedMessageForSession(ctx, reservation.identity, queuedMsg) + case passthroughAttachmentRecovery: + s.restoreQueuedMessage(ctx, queuedMsg) case reservation != nil && reservation.identity.SessionIncarnationID != "" && manualRecovery: s.restoreQueuedMessageForSession(ctx, reservation.identity, queuedMsg) case reservation != nil && reservation.identity.SessionIncarnationID != "": _ = s.requeueMessageForSession(ctx, reservation.identity, queuedMsg, "workflow-auto-start-retry") + case manualRecovery && lifecyclePrompt: + if s.requeueLifecycleMessage(ctx, queuedMsg, queuedMsg.QueuedBy, messageCoalesceKey(queuedMsg)) { + // A reserved lifecycle row remains in the queue while its retry + // successor is inserted. Settle the original reservation so + // manual recovery leaves one visible entry, not two. + s.acknowledgeLifecycleQueueEntry(ctx, queuedMsg.SessionID, queuedMsg) + } case manualRecovery: s.restoreQueuedMessage(ctx, queuedMsg) default: @@ -1449,7 +1531,8 @@ func (s *Service) acknowledgeLifecycleQueueEntry( if s.messageQueue == nil || queuedMsg == nil || !queuedMsg.IsDurableLifecycle() { return } - if err := s.messageQueue.AcknowledgeQueued(ctx, sessionID, queuedMsg.ID); err != nil { + ackCtx := context.WithoutCancel(ctx) + if err := s.messageQueue.AcknowledgeQueued(ackCtx, queuedMsg); err != nil { s.logger.Error("failed to acknowledge accepted lifecycle message", zap.String("session_id", sessionID), zap.String("task_id", queuedMsg.TaskID), @@ -1457,7 +1540,43 @@ func (s *Service) acknowledgeLifecycleQueueEntry( zap.Error(err)) return } - s.publishQueueStatusEvent(ctx, sessionID) + s.publishQueueStatusEvent(ackCtx, sessionID) +} + +func (s *Service) acknowledgeOrdinaryQueueEntry( + ctx context.Context, + sessionID string, + queuedMsg *messagequeue.QueuedMessage, +) { + if s.messageQueue == nil || queuedMsg == nil || queuedMsg.IsDurableLifecycle() { + return + } + ackCtx := context.WithoutCancel(ctx) + if err := s.retrySendNowClaimMutation(ackCtx, func(recoveryCtx context.Context) error { + return s.messageQueue.AcknowledgeQueued(recoveryCtx, queuedMsg) + }); err != nil { + s.logger.Error("failed to acknowledge accepted queue message", + zap.String("session_id", sessionID), + zap.String("task_id", queuedMsg.TaskID), + zap.String("queue_id", queuedMsg.ID), + zap.Error(err)) + } +} + +func (s *Service) queuedMessageAfterDispatch( + ctx context.Context, + queuedMsg *messagequeue.QueuedMessage, + lifecyclePrompt bool, +) func() error { + if lifecyclePrompt || queuedMsg == nil || s.messageQueue == nil || + !s.messageQueue.PendingQueueDispatchPersistenceAvailable() { + return nil + } + return func() error { + return s.retrySendNowClaimMutation(ctx, func(recoveryCtx context.Context) error { + return s.messageQueue.MarkPendingQueueDispatchAccepted(recoveryCtx, queuedMsg) + }) + } } func (s *Service) acknowledgeLifecycleQueueEntryForSession( @@ -1476,7 +1595,10 @@ func (s *Service) acknowledgeLifecycleQueueEntryForSession( } func (s *Service) restoreQueuedMessage(ctx context.Context, queuedMsg *messagequeue.QueuedMessage) { - restored, err := s.messageQueue.RestoreMessage(ctx, queuedMsg) + // Restoration follows a completed take and must survive cancellation of + // the request that initiated the dispatch. + restoreCtx := context.WithoutCancel(ctx) + restored, err := s.messageQueue.RestoreMessage(restoreCtx, queuedMsg) if err != nil { s.logger.Error("failed to restore queued message", zap.String("session_id", queuedMsg.SessionID), @@ -1490,7 +1612,7 @@ func (s *Service) restoreQueuedMessage(ctx context.Context, queuedMsg *messagequ zap.String("task_id", restored.TaskID), zap.String("queue_id", restored.ID), zap.Int64("position", restored.Position)) - s.publishQueueStatusEvent(ctx, queuedMsg.SessionID) + s.publishQueueStatusEvent(restoreCtx, queuedMsg.SessionID) } func (s *Service) restoreQueuedMessageForSession( diff --git a/apps/backend/internal/orchestrator/event_handlers_clarification.go b/apps/backend/internal/orchestrator/event_handlers_clarification.go index 0ae7b58c5b..ad049f42ed 100644 --- a/apps/backend/internal/orchestrator/event_handlers_clarification.go +++ b/apps/backend/internal/orchestrator/event_handlers_clarification.go @@ -595,6 +595,9 @@ func (s *Service) dispatchClarificationResumeLocked(ctx context.Context, data cl if err != nil { return fmt.Errorf("queue clarification resume prompt: %w", err) } + // Queue insertion is observable even when a concurrent drain prevents the + // targeted take, so publish before attempting dispatch. + s.publishQueueStatusEvent(ctx, data.SessionID) dispatched, err := s.takeAndDispatchEntryLocked(ctx, identity, queued.ID) if err != nil { return fmt.Errorf("dispatch clarification resume prompt: %w", err) diff --git a/apps/backend/internal/orchestrator/event_handlers_clarification_test.go b/apps/backend/internal/orchestrator/event_handlers_clarification_test.go index f5ad9aff45..f3fb2a912a 100644 --- a/apps/backend/internal/orchestrator/event_handlers_clarification_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_clarification_test.go @@ -12,6 +12,7 @@ import ( "github.com/kandev/kandev/internal/agent/runtime/lifecycle" "github.com/kandev/kandev/internal/clarification" + "github.com/kandev/kandev/internal/events" "github.com/kandev/kandev/internal/events/bus" "github.com/kandev/kandev/internal/orchestrator/executor" "github.com/kandev/kandev/internal/task/models" @@ -1915,17 +1916,28 @@ func TestDispatchClarificationResumeLocked_ReturnPaths(t *testing.T) { } }) - t.Run("immediate dispatch returns nil", func(t *testing.T) { + t.Run("immediate dispatch publishes queue status", func(t *testing.T) { repo := setupTestRepo(t) agentMgr := &mockAgentManager{isAgentRunning: true, repoForExecutionLookup: repo} svc := createTestServiceWithAgent(repo, newMockStepGetter(), newMockTaskRepo(), agentMgr) svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{}) + recorded := &recordingEventBus{} + svc.eventBus = recorded seedTaskAndSession(t, repo, "t1", "s1", models.TaskSessionStateWaitingForInput) seedExecutorRunning(t, repo, "s1", "t1", "exec-1") if err := svc.dispatchClarificationResumeLocked(ctx, data, "answer"); err != nil { t.Fatalf("expected nil on immediate dispatch, got %v", err) } + var queueStatusEvents int + for _, event := range recorded.events { + if event.subject == events.MessageQueueStatusChanged { + queueStatusEvents++ + } + } + if queueStatusEvents == 0 { + t.Fatal("clarification resume did not publish queue status") + } }) } diff --git a/apps/backend/internal/orchestrator/event_handlers_github_pr_automation_test.go b/apps/backend/internal/orchestrator/event_handlers_github_pr_automation_test.go index 536837b45d..5304f40b3f 100644 --- a/apps/backend/internal/orchestrator/event_handlers_github_pr_automation_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_github_pr_automation_test.go @@ -41,16 +41,16 @@ func (r *lifecycleQueueFailureRepository) InsertOrReplaceLifecycleByCoalesceKeyF // lifecycleAcknowledgingRepository exposes the durable lifecycle completion // boundary. ReserveHead intentionally retains lifecycle entries until the -// executor accepts the prompt and AcknowledgeByID commits the removal. +// executor accepts the prompt and AcknowledgeReserved commits the removal. type lifecycleAcknowledgingRepository struct { messagequeue.Repository acknowledged chan struct{} } -func (r *lifecycleAcknowledgingRepository) AcknowledgeByID( - ctx context.Context, sessionID, entryID string, +func (r *lifecycleAcknowledgingRepository) AcknowledgeReserved( + ctx context.Context, msg *messagequeue.QueuedMessage, ) error { - if err := r.Repository.AcknowledgeByID(ctx, sessionID, entryID); err != nil { + if err := r.Repository.AcknowledgeReserved(ctx, msg); err != nil { return err } close(r.acknowledged) diff --git a/apps/backend/internal/orchestrator/event_handlers_passthrough_running_test.go b/apps/backend/internal/orchestrator/event_handlers_passthrough_running_test.go index 4f3a37a4d9..3b8ed5e036 100644 --- a/apps/backend/internal/orchestrator/event_handlers_passthrough_running_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_passthrough_running_test.go @@ -124,5 +124,8 @@ func TestHandleAgentReady_PassthroughQueuedMessagePublishesAfterWriteFailure(t * require.Equal(t, 1, runningEventCalls, "the running event must still publish after the PTY write fails") require.Empty(t, baseManager.markPassthroughCalls) require.Len(t, baseManager.passthroughStdinCalls, 1) - require.Equal(t, 1, svc.messageQueue.GetStatus(ctx, "s1").Count) + status := svc.messageQueue.GetStatus(ctx, "s1") + require.Len(t, status.Entries, 1, "failed PTY delivery must restore the queued prompt") + require.Equal(t, "queued prompt", status.Entries[0].Content) + require.Equal(t, 1, status.Count) } diff --git a/apps/backend/internal/orchestrator/event_handlers_queue_cancellation_test.go b/apps/backend/internal/orchestrator/event_handlers_queue_cancellation_test.go new file mode 100644 index 0000000000..a84a91e887 --- /dev/null +++ b/apps/backend/internal/orchestrator/event_handlers_queue_cancellation_test.go @@ -0,0 +1,83 @@ +package orchestrator + +import ( + "context" + "testing" + + "github.com/kandev/kandev/internal/events" + "github.com/kandev/kandev/internal/task/models" +) + +func TestRestoreQueuedMessageSurvivesCancelledDispatchContext(t *testing.T) { + repo := setupTestRepo(t) + seedSession(t, repo, "task-restore", "session-restore", "step1") + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + recorded := &recordingEventBus{} + svc.eventBus = recorded + + queued, err := svc.messageQueue.QueueMessage(context.Background(), "session-restore", "task-restore", "prompt", "", "user", false, nil) + if err != nil { + t.Fatalf("queue prompt: %v", err) + } + taken, ok := svc.messageQueue.TakeQueued(context.Background(), "session-restore") + if !ok || taken.ID != queued.ID { + t.Fatalf("take queued message = %#v, %v", taken, ok) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + svc.restoreQueuedMessage(ctx, taken) + + if got := svc.messageQueue.GetStatus(context.Background(), "session-restore").Count; got != 1 { + t.Fatalf("restored queue count = %d, want 1", got) + } + if len(recorded.events) != 1 || recorded.events[0].subject != events.MessageQueueStatusChanged { + t.Fatalf("queue restore status events = %d, want one queue status event", len(recorded.events)) + } + if err := recorded.events[0].ctx.Err(); err != nil { + t.Fatalf("queue restore status context error = %v, want nil", err) + } +} + +func TestTransferQueuedSessionStatePublishesBothSessionStatuses(t *testing.T) { + repo := setupTestRepo(t) + seedTaskAndSession(t, repo, "task-transfer", "session-old", models.TaskSessionStateRunning) + session, err := repo.GetTaskSession(context.Background(), "session-old") + if err != nil { + t.Fatalf("load old session: %v", err) + } + session.ID = "session-new" + if err := repo.CreateTaskSession(context.Background(), session); err != nil { + t.Fatalf("create new session: %v", err) + } + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + recorded := &recordingEventBus{} + svc.eventBus = recorded + + if _, err := svc.messageQueue.QueueMessage(context.Background(), "session-old", "task-transfer", "prompt", "", "user", false, nil); err != nil { + t.Fatalf("queue prompt: %v", err) + } + if err := svc.transferQueuedSessionState(context.Background(), "task-transfer", "session-old", "session-new"); err != nil { + t.Fatalf("transfer queued state: %v", err) + } + + if got := svc.messageQueue.GetStatus(context.Background(), "session-old").Count; got != 0 { + t.Fatalf("old queue count = %d, want 0", got) + } + if got := svc.messageQueue.GetStatus(context.Background(), "session-new").Count; got != 1 { + t.Fatalf("new queue count = %d, want 1", got) + } + if len(recorded.events) != 2 { + t.Fatalf("queue transfer status events = %d, want 2", len(recorded.events)) + } + seen := map[string]bool{} + for _, event := range recorded.events { + if event.subject != events.MessageQueueStatusChanged { + t.Fatalf("queue transfer event subject = %q, want %q", event.subject, events.MessageQueueStatusChanged) + } + seen[event.event.Data.(map[string]interface{})["session_id"].(string)] = true + } + if !seen["session-old"] || !seen["session-new"] { + t.Fatalf("queue transfer status sessions = %#v, want both sessions", seen) + } +} diff --git a/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go b/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go index 908e4a689f..8a23693ae3 100644 --- a/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go @@ -125,11 +125,17 @@ func TestExecuteQueuedMessage_LifecycleReselectsReplacementSession(t *testing.T) agentMgr := &mockAgentManager{isAgentRunning: true, repoForExecutionLookup: repo} svc := createTestServiceWithAgent(repo, newMockStepGetter(), newMockTaskRepo(), agentMgr) svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{}) - queued := &messagequeue.QueuedMessage{ - ID: "lifecycle-reselect", SessionID: "s1", TaskID: "t1", Content: "merged lifecycle prompt", - Metadata: map[string]interface{}{ - "origin": githubPRAutomationOrigin, messagequeue.MetadataCoalesceKey: "github-pr:repo:1:merged", - }, + _, _, accepted, err := svc.messageQueue.QueueLifecycleMessageWithCoalesceKey( + ctx, "s1", "t1", "merged lifecycle prompt", "", messagequeue.QueuedByWorkflow, + false, nil, map[string]interface{}{"origin": githubPRAutomationOrigin}, + "github-pr:repo:1:merged", true, + ) + if err != nil || !accepted { + t.Fatalf("queue lifecycle prompt: accepted=%t err=%v", accepted, err) + } + queued, ok := svc.messageQueue.ReserveQueued(ctx, "s1") + if !ok { + t.Fatal("reserve lifecycle prompt") } // The event was accepted for s1, then s2 becomes the current primary at @@ -145,6 +151,9 @@ func TestExecuteQueuedMessage_LifecycleReselectsReplacementSession(t *testing.T) svc.markQueuedDispatchInFlight("s1", queued.ID) svc.executeQueuedMessage("s1", queued) + if stale, ok := svc.messageQueue.ReserveQueued(ctx, "s1"); ok { + t.Fatalf("old lifecycle reservation remained after reselection: %#v", stale) + } status := svc.messageQueue.GetStatus(ctx, "s2") if status.Count != 1 { t.Fatalf("replacement lifecycle retries = %d, want 1", status.Count) diff --git a/apps/backend/internal/orchestrator/event_handlers_queue_passthrough_attachment_test.go b/apps/backend/internal/orchestrator/event_handlers_queue_passthrough_attachment_test.go new file mode 100644 index 0000000000..c355c4f2d4 --- /dev/null +++ b/apps/backend/internal/orchestrator/event_handlers_queue_passthrough_attachment_test.go @@ -0,0 +1,117 @@ +package orchestrator + +import ( + "context" + "strings" + "testing" + + "github.com/kandev/kandev/internal/orchestrator/executor" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/orchestrator/watcher" + wfmodels "github.com/kandev/kandev/internal/workflow/models" + v1 "github.com/kandev/kandev/pkg/api/v1" +) + +func TestHandleAgentReady_PassthroughAttachmentOnlyOrdinaryMessageUsesAttachmentAwarePrompt(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-attachment-only", "session-attachment-only", "step1") + session, err := repo.GetTaskSession(ctx, "session-attachment-only") + if err != nil { + t.Fatalf("load session: %v", err) + } + session.WorkspacePath = t.TempDir() + if err := repo.UpdateTaskSession(ctx, session); err != nil { + t.Fatalf("update session workspace: %v", err) + } + seedExecutorRunning(t, repo, "session-attachment-only", "task-attachment-only", "execution-attachment-only") + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "task-attachment-only", v1.TaskStateReview) + agentMgr := &mockAgentManager{ + isPassthrough: true, + isAgentRunning: true, + repoForExecutionLookup: repo, + } + steps := newMockStepGetter() + steps.steps["step1"] = &wfmodels.WorkflowStep{ID: "step1", WorkflowID: "workflow1"} + svc := createTestServiceWithAgent(repo, steps, taskRepo, agentMgr) + svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{}) + + _, err = svc.messageQueue.QueueMessage( + ctx, + "session-attachment-only", + "task-attachment-only", + "", + "", + messagequeue.QueuedByUser, + false, + []messagequeue.MessageAttachment{{Type: "image", Data: "AQ==", MimeType: "image/png"}}, + ) + if err != nil { + t.Fatalf("queue attachment-only message: %v", err) + } + + svc.handleAgentReady(ctx, watcher.AgentEventData{ + TaskID: "task-attachment-only", + SessionID: "session-attachment-only", + }) + + if len(agentMgr.passthroughStdinCalls) != 1 { + t.Fatalf("passthrough writes = %+v, want one attachment-aware write", agentMgr.passthroughStdinCalls) + } + if !strings.Contains(agentMgr.passthroughStdinCalls[0].Data, ".kandev/attachments") { + t.Fatalf("passthrough write omitted attachment path: %q", agentMgr.passthroughStdinCalls[0].Data) + } +} + +func TestHandleAgentReady_PassthroughQueuedAttachmentUsesAttachmentAwarePrompt(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-attachment", "session-attachment", "step1") + session, err := repo.GetTaskSession(ctx, "session-attachment") + if err != nil { + t.Fatalf("load session: %v", err) + } + session.WorkspacePath = t.TempDir() + if err := repo.UpdateTaskSession(ctx, session); err != nil { + t.Fatalf("update session workspace: %v", err) + } + seedExecutorRunning(t, repo, "session-attachment", "task-attachment", "execution-attachment") + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "task-attachment", v1.TaskStateReview) + agentMgr := &mockAgentManager{ + isPassthrough: true, + isAgentRunning: true, + repoForExecutionLookup: repo, + } + steps := newMockStepGetter() + steps.steps["step1"] = &wfmodels.WorkflowStep{ID: "step1", WorkflowID: "workflow1"} + svc := createTestServiceWithAgent(repo, steps, taskRepo, agentMgr) + svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{}) + + _, err = svc.messageQueue.QueueMessage( + ctx, + "session-attachment", + "task-attachment", + "read attached file", + "", + messagequeue.QueuedByUser, + false, + []messagequeue.MessageAttachment{{Type: "resource", Data: "AQ==", MimeType: "text/plain", Name: "note.txt"}}, + ) + if err != nil { + t.Fatalf("queue message with attachment: %v", err) + } + + svc.handleAgentReady(ctx, watcher.AgentEventData{ + TaskID: "task-attachment", + SessionID: "session-attachment", + }) + + if len(agentMgr.passthroughStdinCalls) != 1 { + t.Fatalf("passthrough writes = %+v, want one attachment-aware write", agentMgr.passthroughStdinCalls) + } + if !strings.Contains(agentMgr.passthroughStdinCalls[0].Data, ".kandev/attachments") { + t.Fatalf("passthrough write omitted attachment path: %q", agentMgr.passthroughStdinCalls[0].Data) + } +} diff --git a/apps/backend/internal/orchestrator/event_handlers_queue_recovery_test.go b/apps/backend/internal/orchestrator/event_handlers_queue_recovery_test.go new file mode 100644 index 0000000000..a17aeff502 --- /dev/null +++ b/apps/backend/internal/orchestrator/event_handlers_queue_recovery_test.go @@ -0,0 +1,184 @@ +package orchestrator + +import ( + "context" + "database/sql" + "errors" + "path/filepath" + "testing" + + "github.com/jmoiron/sqlx" + "github.com/kandev/kandev/internal/orchestrator/executor" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/orchestrator/watcher" + wfmodels "github.com/kandev/kandev/internal/workflow/models" + v1 "github.com/kandev/kandev/pkg/api/v1" + _ "github.com/mattn/go-sqlite3" +) + +func TestDispatchTakenQueuedMessageDiscardsEmptyOrdinaryEntry(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-empty", "session-empty", "step1") + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + + _, err := svc.messageQueue.QueueMessage(ctx, "session-empty", "task-empty", "", "", messagequeue.QueuedByUser, false, nil) + if err != nil { + t.Fatalf("queue empty ordinary entry: %v", err) + } + queued, ok := svc.messageQueue.TakeQueued(ctx, "session-empty") + if !ok || queued == nil { + t.Fatal("take empty ordinary entry") + } + + if svc.dispatchTakenQueuedMessage(ctx, "session-empty", queued, true) { + t.Fatal("empty ordinary entry was dispatched") + } + if got := svc.messageQueue.GetStatus(ctx, "session-empty").Count; got != 0 { + t.Fatalf("empty ordinary entry count = %d, want discarded entry", got) + } +} + +func TestDispatchTakenQueuedMessageAcknowledgesEmptyLifecycleEntry(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + _, _, accepted, err := svc.messageQueue.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-empty-lifecycle", "task-empty-lifecycle", "", "", messagequeue.QueuedByWorkflow, + false, nil, map[string]interface{}{"origin": githubPRAutomationOrigin}, "empty-lifecycle", true, + ) + if err != nil || !accepted { + t.Fatalf("queue empty lifecycle entry: accepted=%v err=%v", accepted, err) + } + queued, ok := svc.messageQueue.ReserveQueued(ctx, "session-empty-lifecycle") + if !ok || queued == nil { + t.Fatal("reserve empty lifecycle entry") + } + + if svc.dispatchTakenQueuedMessage(ctx, "session-empty-lifecycle", queued, true) { + t.Fatal("empty lifecycle entry was dispatched") + } + entries, _, err := svc.messageQueue.SnapshotSession(ctx, "session-empty-lifecycle") + if err != nil { + t.Fatalf("snapshot empty lifecycle queue: %v", err) + } + if len(entries) != 0 { + t.Fatalf("empty lifecycle entries after discard = %d, want 0", len(entries)) + } +} +func TestDispatchTakenQueuedMessageAcknowledgesEmptyLifecycleAfterCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + repo := setupTestRepo(t) + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + rawQueueDB, err := sql.Open("sqlite3", filepath.Join(t.TempDir(), "queue.db")) + if err != nil { + t.Fatalf("open queue database: %v", err) + } + rawQueueDB.SetMaxOpenConns(1) + rawQueueDB.SetMaxIdleConns(1) + queueDB := sqlx.NewDb(rawQueueDB, "sqlite3") + t.Cleanup(func() { _ = queueDB.Close() }) + if _, err := queueDB.Exec(`CREATE TABLE tasks (id TEXT PRIMARY KEY, archived_at TIMESTAMP NULL, updated_at TIMESTAMP NOT NULL)`); err != nil { + t.Fatalf("create queue task stub: %v", err) + } + if _, err := queueDB.Exec(`INSERT INTO tasks (id, updated_at) VALUES (?, CURRENT_TIMESTAMP)`, "task-empty-lifecycle-cancelled"); err != nil { + t.Fatalf("seed queue task stub: %v", err) + } + queueRepo, err := messagequeue.NewSQLiteRepository(queueDB, queueDB) + if err != nil { + t.Fatalf("create SQLite queue repository: %v", err) + } + svc.messageQueue = messagequeue.NewService(queueRepo, messagequeue.DefaultMaxPerSession, testLogger()) + _, _, accepted, err := svc.messageQueue.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-empty-lifecycle-cancelled", "task-empty-lifecycle-cancelled", "", "", + messagequeue.QueuedByWorkflow, false, nil, + map[string]interface{}{"origin": githubPRAutomationOrigin}, "empty-lifecycle-cancelled", true, + ) + if err != nil || !accepted { + t.Fatalf("queue empty lifecycle entry: accepted=%v err=%v", accepted, err) + } + queued, ok := svc.messageQueue.ReserveQueued(ctx, "session-empty-lifecycle-cancelled") + if !ok || queued == nil { + t.Fatal("reserve empty lifecycle entry") + } + cancel() + + if svc.dispatchTakenQueuedMessage(ctx, queued.SessionID, queued, true) { + t.Fatal("empty lifecycle entry was dispatched") + } + var durableCount int + if err := queueDB.GetContext(context.Background(), &durableCount, `SELECT COUNT(*) FROM queued_messages WHERE session_id = ?`, queued.SessionID); err != nil { + t.Fatalf("count durable lifecycle entries: %v", err) + } + if durableCount != 0 { + t.Fatalf("durable lifecycle entries after discard = %d, want 0", durableCount) + } + entries, _, err := svc.messageQueue.SnapshotSession(context.Background(), queued.SessionID) + if err != nil { + t.Fatalf("snapshot empty lifecycle queue: %v", err) + } + if len(entries) != 0 { + t.Fatalf("empty lifecycle entries after discard = %d, want 0", len(entries)) + } +} + +func TestTakeAndMergeHandoffMessageDiscardsEmptyEntry(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-empty-handoff", "session-empty-handoff", "step1") + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + _, err := svc.messageQueue.QueueMessage(ctx, "session-empty-handoff", "task-empty-handoff", "", "", messagequeue.QueuedByUser, false, nil) + if err != nil { + t.Fatalf("queue empty handoff: %v", err) + } + + msg, prompt, attachments, references, handoff := svc.takeAndMergeHandoffMessage(ctx, "session-empty-handoff", "base") + if msg != nil || prompt != "base" || attachments != nil || references != nil || handoff != "" { + t.Fatalf("empty handoff result = msg:%v prompt:%q attachments:%v references:%v handoff:%q", msg, prompt, attachments, references, handoff) + } + if got := svc.messageQueue.GetStatus(ctx, "session-empty-handoff").Count; got != 0 { + t.Fatalf("empty handoff count = %d, want discarded entry", got) + } +} + +func TestHandleAgentReadyPassthroughAttachmentFailureRestoresEntry(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-attachment-failure", "session-attachment-failure", "step1") + session, err := repo.GetTaskSession(ctx, "session-attachment-failure") + if err != nil { + t.Fatalf("load session: %v", err) + } + session.WorkspacePath = t.TempDir() + if err := repo.UpdateTaskSession(ctx, session); err != nil { + t.Fatalf("update session workspace: %v", err) + } + seedExecutorRunning(t, repo, "session-attachment-failure", "task-attachment-failure", "execution-attachment-failure") + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "task-attachment-failure", v1.TaskStateReview) + agentMgr := &mockAgentManager{ + isPassthrough: true, + isAgentRunning: true, + passthroughStdinErr: errors.New("PTY write failed"), + repoForExecutionLookup: repo, + } + steps := newMockStepGetter() + steps.steps["step1"] = &wfmodels.WorkflowStep{ID: "step1", WorkflowID: "workflow1"} + svc := createTestServiceWithAgent(repo, steps, taskRepo, agentMgr) + svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{}) + + _, err = svc.messageQueue.QueueMessage( + ctx, "session-attachment-failure", "task-attachment-failure", "read attached file", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{Type: "resource", Data: "AQ==", MimeType: "text/plain", Name: "note.txt"}}, + ) + if err != nil { + t.Fatalf("queue attachment entry: %v", err) + } + + svc.handleAgentReady(ctx, watcher.AgentEventData{TaskID: "task-attachment-failure", SessionID: "session-attachment-failure"}) + if got := svc.messageQueue.GetStatus(ctx, "session-attachment-failure").Count; got != 1 { + t.Fatalf("attachment entry count after PTY failure = %d, want restored entry", got) + } +} diff --git a/apps/backend/internal/orchestrator/event_handlers_queue_round64_test.go b/apps/backend/internal/orchestrator/event_handlers_queue_round64_test.go new file mode 100644 index 0000000000..59986612db --- /dev/null +++ b/apps/backend/internal/orchestrator/event_handlers_queue_round64_test.go @@ -0,0 +1,198 @@ +package orchestrator + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/task/models" +) + +type lifecycleRetryRoutingRepository struct { + messagequeue.Repository + lifecycleRetries int + ordinaryRetries int +} + +func (r *lifecycleRetryRoutingRepository) RequeuePreservingFIFO(ctx context.Context, msg *messagequeue.QueuedMessage) error { + r.ordinaryRetries++ + return r.Repository.RequeuePreservingFIFO(ctx, msg) +} + +func (r *lifecycleRetryRoutingRepository) InsertOrReplaceLifecycleByCoalesceKey( + ctx context.Context, + msg *messagequeue.QueuedMessage, + coalesceKey string, + maxPerSession int, + allowInsert bool, +) (*messagequeue.QueuedMessage, bool, error) { + r.lifecycleRetries++ + return r.Repository.InsertOrReplaceLifecycleByCoalesceKey(ctx, msg, coalesceKey, maxPerSession, allowInsert) +} + +func TestRequeueDurableLifecycleMarkerWithoutOriginUsesLifecyclePath(t *testing.T) { + ctx := context.Background() + repo := &lifecycleRetryRoutingRepository{Repository: messagequeue.NewMemoryRepository()} + queue := messagequeue.NewService(repo, messagequeue.DefaultMaxPerSession, testLogger()) + svc := &Service{messageQueue: queue, logger: testLogger()} + + queued, _, accepted, err := queue.QueueLifecycleMessageWithCoalesceKey( + ctx, + "session-round64", + "task-round64", + "durable prompt", + "", + messagequeue.QueuedByWorkflow, + false, + nil, + map[string]interface{}{}, + "lifecycle:round64", + true, + ) + if err != nil || !accepted || queued == nil { + t.Fatalf("queue durable lifecycle message: accepted=%t entry=%v err=%v", accepted, queued, err) + } + + reserved, dispatched, autoRun := queue.ReserveQueuedWithAutoRun(ctx, "session-round64") + if !autoRun || !dispatched || reserved == nil { + t.Fatalf("reserve durable lifecycle message: auto_run=%t dispatched=%t message=%v", autoRun, dispatched, reserved) + } + if !reserved.IsReservedLifecycleDelivery() { + t.Fatal("reserved lifecycle message did not carry delivery reservation") + } + + // Legacy durable rows can have the durable marker without a provider origin. + // Their retry must retain lifecycle generation and inactive-task fencing, + // rather than taking the ordinary user FIFO retry path. + repo.lifecycleRetries = 0 + repo.ordinaryRetries = 0 + svc.requeueMessage(ctx, reserved, reserved.QueuedBy) + + if repo.lifecycleRetries != 1 { + t.Fatalf("lifecycle retry repository calls = %d, want 1", repo.lifecycleRetries) + } + if repo.ordinaryRetries != 0 { + t.Fatalf("ordinary retry repository calls = %d, want 0", repo.ordinaryRetries) + } +} + +func TestLifecycleManualRecoveryReplacesReservedRow(t *testing.T) { + ctx := context.Background() + repo := messagequeue.NewMemoryRepository() + queue := messagequeue.NewService(repo, messagequeue.DefaultMaxPerSession, testLogger()) + svc := &Service{messageQueue: queue, logger: testLogger()} + + queued, _, accepted, err := queue.QueueLifecycleMessageWithCoalesceKey( + ctx, + "session-round64-recovery", + "task-round64-recovery", + "durable prompt", + "", + messagequeue.QueuedByWorkflow, + false, + nil, + map[string]interface{}{"origin": githubPRAutomationOrigin}, + "lifecycle:round64:recovery", + true, + ) + if err != nil || !accepted || queued == nil { + t.Fatalf("queue durable lifecycle message: accepted=%t entry=%v err=%v", accepted, queued, err) + } + reserved, dispatched, autoRun := queue.ReserveQueuedWithAutoRun(ctx, "session-round64-recovery") + if !autoRun || !dispatched || reserved == nil { + t.Fatalf("reserve durable lifecycle message: auto_run=%t dispatched=%t message=%v", autoRun, dispatched, reserved) + } + + svc.handleQueuedMessageExecutionError( + ctx, + "session-round64-recovery", + reserved, + nil, + true, + false, + errors.New("UsageLimitExceeded"), + ) + + entries, err := repo.ListBySession(ctx, "session-round64-recovery") + if err != nil { + t.Fatalf("list recovered lifecycle message: %v", err) + } + if len(entries) != 1 { + t.Fatalf("recovered lifecycle entries = %d, want exactly one", len(entries)) + } + if entries[0].IsReservedInFlight() { + t.Fatal("recovered lifecycle message remained reserved") + } +} + +func TestReuseSessionRestoresPrimaryWhenQueueTransferFails(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-round64-reuse", "session-round64-current", "step-round64") + current, err := repo.GetTaskSession(ctx, "session-round64-current") + if err != nil { + t.Fatal(err) + } + current.State = models.TaskSessionStateRunning + current.IsPrimary = true + if err := repo.UpdateTaskSession(ctx, current); err != nil { + t.Fatal(err) + } + existing := &models.TaskSession{ + ID: "session-round64-existing", + TaskID: current.TaskID, + State: models.TaskSessionStateWaitingForInput, + AgentProfileID: "profile-round64", + StartedAt: current.StartedAt.Add(time.Second), + UpdatedAt: current.UpdatedAt.Add(time.Second), + } + if err := repo.CreateTaskSession(ctx, existing); err != nil { + t.Fatal(err) + } + + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + transferErr := errors.New("queue transfer failed") + svc.messageQueue = messagequeue.NewService( + &transferFailureQueueRepository{Repository: messagequeue.NewMemoryRepository(), err: transferErr}, + messagequeue.DefaultMaxPerSession, + testLogger(), + ) + if _, err := svc.messageQueue.QueueMessage( + ctx, + current.ID, + current.TaskID, + "handoff", + "", + messagequeue.QueuedByUser, + false, + nil, + ); err != nil { + t.Fatal(err) + } + + _, err = svc.reuseSessionForStepWithEndPolicy( + ctx, current.TaskID, current, existing, models.WorkflowProfileSessionEndPolicyComplete, + ) + if !errors.Is(err, transferErr) { + t.Fatalf("reuseSessionForStep error = %v, want queue transfer failure", err) + } + + sessions, err := repo.ListTaskSessions(ctx, current.TaskID) + if err != nil { + t.Fatal(err) + } + for _, session := range sessions { + switch session.ID { + case current.ID: + if !session.IsPrimary || session.State != models.TaskSessionStateRunning { + t.Fatalf("current session after failed transfer = primary %t state %q, want primary running", session.IsPrimary, session.State) + } + case existing.ID: + if session.IsPrimary { + t.Fatal("reused session remained primary after failed queue transfer") + } + } + } +} diff --git a/apps/backend/internal/orchestrator/event_handlers_queue_status_test.go b/apps/backend/internal/orchestrator/event_handlers_queue_status_test.go new file mode 100644 index 0000000000..14b9329eab --- /dev/null +++ b/apps/backend/internal/orchestrator/event_handlers_queue_status_test.go @@ -0,0 +1,76 @@ +package orchestrator + +import ( + "context" + "testing" + + "github.com/kandev/kandev/internal/events" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/orchestrator/watcher" + "github.com/kandev/kandev/internal/workflow/models" + v1 "github.com/kandev/kandev/pkg/api/v1" +) + +func TestHandleAgentReadyPassthroughPublishesStatusAfterReservation(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-passthrough-status", "session-passthrough-status", "step1") + agentMgr := &mockAgentManager{ + isPassthrough: true, + isAgentRunning: true, + repoForExecutionLookup: repo, + } + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "task-passthrough-status", v1.TaskStateReview) + steps := newMockStepGetter() + steps.steps["step1"] = &models.WorkflowStep{ID: "step1", WorkflowID: "wf1"} + svc := createTestServiceWithAgent(repo, steps, taskRepo, agentMgr) + recorded := &recordingEventBus{} + svc.eventBus = recorded + + if _, err := svc.messageQueue.QueueMessage( + ctx, + "session-passthrough-status", + "task-passthrough-status", + "prompt", + "", + messagequeue.QueuedByUser, + false, + nil, + ); err != nil { + t.Fatalf("queue prompt: %v", err) + } + + svc.handleAgentReady(ctx, watcher.AgentEventData{ + TaskID: "task-passthrough-status", + SessionID: "session-passthrough-status", + }) + + var queueStatusEvents int + for _, event := range recorded.events { + if event.subject == events.MessageQueueStatusChanged { + queueStatusEvents++ + } + } + if queueStatusEvents != 1 { + t.Fatalf("queue status events after passthrough reservation = %d, want 1", queueStatusEvents) + } +} +func TestPublishQueueStatusDetachesCancelledRequest(t *testing.T) { + repo := setupTestRepo(t) + seedSession(t, repo, "task-cancelled-status", "session-cancelled-status", "step1") + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + recorded := &recordingEventBus{} + svc.eventBus = recorded + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + svc.PublishQueueStatusEvent(ctx, "session-cancelled-status") + + if len(recorded.events) != 1 { + t.Fatalf("queue status events = %d, want 1", len(recorded.events)) + } + if err := recorded.events[0].ctx.Err(); err != nil { + t.Fatalf("queue status event context error = %v, want nil", err) + } +} diff --git a/apps/backend/internal/orchestrator/event_handlers_streaming_test.go b/apps/backend/internal/orchestrator/event_handlers_streaming_test.go index 3a198dd4a5..0b345d9a43 100644 --- a/apps/backend/internal/orchestrator/event_handlers_streaming_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_streaming_test.go @@ -31,6 +31,7 @@ type recordingEventBus struct { type recordedEvent struct { subject string event *bus.Event + ctx context.Context } // taskServiceStateRepository exercises runtime task-state reconciliation @@ -341,8 +342,8 @@ func (r failSetSessionMetadataRepo) SetSessionMetadataKey( return errors.New("set session metadata failed") } -func (b *recordingEventBus) Publish(_ context.Context, subject string, event *bus.Event) error { - b.events = append(b.events, recordedEvent{subject: subject, event: event}) +func (b *recordingEventBus) Publish(ctx context.Context, subject string, event *bus.Event) error { + b.events = append(b.events, recordedEvent{subject: subject, event: event, ctx: ctx}) return nil } func (b *recordingEventBus) Subscribe(string, bus.EventHandler) (bus.Subscription, error) { diff --git a/apps/backend/internal/orchestrator/event_handlers_workflow.go b/apps/backend/internal/orchestrator/event_handlers_workflow.go index 02ba693ea8..5a251bd264 100644 --- a/apps/backend/internal/orchestrator/event_handlers_workflow.go +++ b/apps/backend/internal/orchestrator/event_handlers_workflow.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "reflect" "strconv" "strings" "sync" @@ -38,8 +39,22 @@ var ( errReusableSessionNoLongerActive = errors.New("reusable session is no longer active") errWorkflowAutoStartSessionTerminalized = errors.New("workflow auto-start session terminalized") errContextResetCancellationConflict = errors.New("context reset cancellation is already in progress") + errSessionAttachmentTransferUnavailable = errors.New("session attachment transfer service is unavailable") ) +func sessionAttachmentTransfererAvailable(transfer SessionAttachmentTransferer) bool { + if transfer == nil { + return false + } + value := reflect.ValueOf(transfer) + switch value.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return !value.IsNil() + default: + return true + } +} + type workflowAutoStartSessionTerminalizedError struct { state models.TaskSessionState } @@ -2532,6 +2547,224 @@ func (s *Service) findReusableSessionForProfile(ctx context.Context, taskID, pro return best, nil } +// transferQueuedSessionState keeps queue rows and their claimed attachment +// bindings aligned when workflow session ownership changes. +func (s *Service) transferQueuedSessionState(ctx context.Context, taskID, oldSessionID, newSessionID string) error { + if s.messageQueue == nil { + return nil + } + transferErr := s.messageQueue.TransferSessionWithDurableAttachmentPreparation( + ctx, + taskID, + oldSessionID, + newSessionID, + func(admittedCtx context.Context, attachmentIDs []string) error { + if len(attachmentIDs) == 0 { + return nil + } + if !sessionAttachmentTransfererAvailable(s.sessionAttachmentTransferer) { + return errSessionAttachmentTransferUnavailable + } + return s.sessionAttachmentTransferer.TransferSessionMessageAttachments( + admittedCtx, taskID, oldSessionID, newSessionID, attachmentIDs, + ) + }, + func(rollbackCtx context.Context, attachmentIDs []string) error { + if len(attachmentIDs) == 0 { + return nil + } + if !sessionAttachmentTransfererAvailable(s.sessionAttachmentTransferer) { + return errSessionAttachmentTransferUnavailable + } + return s.sessionAttachmentTransferer.TransferSessionMessageAttachments( + rollbackCtx, taskID, newSessionID, oldSessionID, attachmentIDs, + ) + }, + ) + if transferErr != nil { + return fmt.Errorf("transfer queued state: %w", transferErr) + } + s.publishQueueStatusEvent(ctx, oldSessionID) + s.publishQueueStatusEvent(ctx, newSessionID) + return nil +} + +func (s *Service) reconcileSessionTransferCompensationsOnStartup(ctx context.Context) error { + if s.messageQueue == nil || !s.messageQueue.SessionTransferCompensationPersistenceAvailable() { + return nil + } + compensations, err := s.messageQueue.ListSessionTransferCompensations(ctx) + if err != nil { + return fmt.Errorf("list session transfer compensations: %w", err) + } + if sessionTransferCompensationsNeedAttachmentTransfer(compensations) && + !sessionAttachmentTransfererAvailable(s.sessionAttachmentTransferer) { + return errors.New("reconcile session transfer compensations: attachment transfer service is unavailable") + } + for _, compensation := range compensations { + if err := s.reconcileSessionTransferCompensation(ctx, compensation); err != nil { + return err + } + } + + return nil +} + +func sessionTransferCompensationsNeedAttachmentTransfer( + compensations []messagequeue.SessionTransferCompensation, +) bool { + for _, compensation := range compensations { + if len(compensation.AttachmentIDs) > 0 { + return true + } + } + return false +} + +func (s *Service) reconcileSessionTransferCompensation( + ctx context.Context, + compensation messagequeue.SessionTransferCompensation, +) (resultErr error) { + ownerID, err := s.messageQueue.ClaimSessionTransferCompensationRecovery(ctx, compensation) + if err != nil { + return fmt.Errorf("claim session transfer compensation recovery: %w", err) + } + transferCtx, cancelTransfer := context.WithCancel(ctx) + defer cancelTransfer() + operationCtx, stopLeaseRenewal := s.messageQueue.MaintainSessionTransferCompensationLeaseWithCancel( + transferCtx, compensation, ownerID, cancelTransfer, + ) + stopped := false + defer func() { + if stopped { + return + } + if leaseErr := stopLeaseRenewal(); leaseErr != nil { + resultErr = errors.Join(resultErr, fmt.Errorf("session transfer lease lost: %w", leaseErr)) + } + }() + targetSessionID, err := s.sessionTransferCompensationTarget(operationCtx, compensation) + if err != nil { + return err + } + fromSessionID, toSessionID := compensation.FromSessionID, compensation.ToSessionID + if targetSessionID == compensation.FromSessionID { + fromSessionID, toSessionID = compensation.ToSessionID, compensation.FromSessionID + } + var transferErr error + if len(compensation.AttachmentIDs) > 0 { + transferErr = s.sessionAttachmentTransferer.TransferSessionMessageAttachments( + operationCtx, + compensation.TaskID, + fromSessionID, + toSessionID, + compensation.AttachmentIDs, + ) + } + leaseErr := stopLeaseRenewal() + stopped = true + if transferErr != nil { + transferErr = fmt.Errorf("reconcile session transfer attachments: %w", transferErr) + if leaseErr != nil { + return errors.Join(transferErr, fmt.Errorf("session transfer lease lost: %w", leaseErr)) + } + return transferErr + } + if leaseErr != nil { + return fmt.Errorf("session transfer lease lost: %w", leaseErr) + } + if err := s.messageQueue.DeleteClaimedSessionTransferCompensation( + context.WithoutCancel(ctx), + compensation, + ownerID, + ); err != nil { + return fmt.Errorf("delete session transfer compensation: %w", err) + } + return nil +} + +func (s *Service) pendingDispatchTransferCompensationTarget( + ctx context.Context, + compensation messagequeue.SessionTransferCompensation, +) (string, bool, error) { + entryIDs := make(map[string]struct{}, len(compensation.EntryIDs)) + for _, entryID := range compensation.EntryIDs { + entryIDs[entryID] = struct{}{} + } + dispatches, err := s.messageQueue.ListPendingQueueDispatches(ctx) + if err != nil { + return "", false, fmt.Errorf("list compensated queue dispatches: %w", err) + } + for _, dispatch := range dispatches { + if _, ok := entryIDs[dispatch.Message.ID]; !ok { + continue + } + sessionID := dispatch.Message.SessionID + if sessionID != compensation.FromSessionID && sessionID != compensation.ToSessionID { + return "", false, fmt.Errorf( + "compensated queue dispatch %s belongs to unexpected session %s", + dispatch.Message.ID, + sessionID, + ) + } + return sessionID, true, nil + } + return "", false, nil +} + +func (s *Service) sessionTransferCompensationTarget( + ctx context.Context, + compensation messagequeue.SessionTransferCompensation, +) (string, error) { + for _, entryID := range compensation.EntryIDs { + entry, err := s.messageQueue.FindEntryByID(ctx, entryID) + if err == nil { + if entry.SessionID != compensation.FromSessionID && entry.SessionID != compensation.ToSessionID { + return "", fmt.Errorf( + "compensated queue entry %s belongs to unexpected session %s", + entryID, + entry.SessionID, + ) + } + return entry.SessionID, nil + } + if !errors.Is(err, messagequeue.ErrEntryNotFound) { + return "", fmt.Errorf("locate compensated queue entry %s: %w", entryID, err) + } + } + if sessionID, found, err := s.pendingDispatchTransferCompensationTarget( + ctx, compensation, + ); err != nil { + return "", err + } else if found { + return sessionID, nil + } + for _, locator := range compensation.CleanupLocators { + cleanup, err := s.messageQueue.GetAttachmentCleanup( + ctx, locator.SessionID, locator.EntryID, locator.OperationID, + ) + if err != nil { + return "", fmt.Errorf("locate compensated attachment cleanup %s: %w", locator.EntryID, err) + } + if cleanup == nil { + continue + } + currentSessionID := cleanup.CurrentSessionID + if currentSessionID == "" { + currentSessionID = cleanup.SessionID + } + if currentSessionID != compensation.FromSessionID && currentSessionID != compensation.ToSessionID { + return "", fmt.Errorf( + "compensated attachment cleanup %s belongs to unexpected session %s", + locator.EntryID, + currentSessionID, + ) + } + return currentSessionID, nil + } + return compensation.ToSessionID, nil +} + func (s *Service) reuseSessionForStepWithEndPolicy( ctx context.Context, taskID string, @@ -2559,7 +2792,6 @@ func (s *Service) reuseSessionForStepWithEndPolicy( } else if task != nil { s.publishTaskUpdated(ctx, task) } - s.tagSessionAsWorkflowSwitched(ctx, existing.ID) // Parking can fail while persisting stop ownership. Do it before queue // transfer so that failure leaves both source and destination queues under @@ -2570,6 +2802,7 @@ func (s *Service) reuseSessionForStepWithEndPolicy( } return existing, nil } + s.tagSessionAsWorkflowSwitched(ctx, existing.ID) if err := s.transferWorkflowProfileSwitchQueue(ctx, currentSession.ID, existing.ID); err != nil { transferErr := fmt.Errorf("transfer queued state to reused session: %w", err) @@ -2648,10 +2881,6 @@ func (s *Service) createNewSessionForStepWithEndPolicy( return nil, err } - // Tag the session as workflow-spawned for provenance: its agent profile - // was selected by the workflow step override rather than direct user choice. - s.tagSessionAsWorkflowSwitched(ctx, newSession.ID) - // Promote the new session to primary so it's loaded when navigating back to this task. // Use SetPrimarySession (not repo.SetSessionPrimary) to broadcast a task.updated WS // event — the frontend reads primarySessionId from the task to render the star icon. @@ -2748,6 +2977,11 @@ func (s *Service) prepareWorkflowReplacementSession( return nil, fmt.Errorf("failed to attach workflow replacement workspace: %w", err) } + newSession, err = s.repo.GetTaskSession(ctx, sessionID) + if err != nil { + return nil, fmt.Errorf("failed to get new session: %w", err) + } + return newSession, nil } @@ -2755,7 +2989,14 @@ func (s *Service) transferWorkflowProfileSwitchQueue(ctx context.Context, fromSe if s.messageQueue == nil { return nil } - if err := s.transferCurrentQueueSessions(ctx, fromSessionID, toSessionID); err != nil { + from, err := s.repo.GetTaskSession(ctx, fromSessionID) + if err != nil { + return fmt.Errorf("load source queue session: %w", err) + } + if from == nil { + return fmt.Errorf("load source queue session %q: not found", fromSessionID) + } + if err := s.transferQueuedSessionState(ctx, from.TaskID, fromSessionID, toSessionID); err != nil { return fmt.Errorf("transfer queue to new workflow session: %w", err) } return nil @@ -2876,42 +3117,6 @@ func (s *Service) restoreWorkflowProfileSwitchSourcePrimary(ctx context.Context, return nil } -func (s *Service) transferCurrentQueueSessions(ctx context.Context, fromSessionID, toSessionID string) error { - from, err := s.repo.GetTaskSession(ctx, fromSessionID) - if err != nil { - return fmt.Errorf("load source queue session: %w", err) - } - if from == nil { - return fmt.Errorf("load source queue session %q: not found", fromSessionID) - } - to, err := s.repo.GetTaskSession(ctx, toSessionID) - if err != nil { - return fmt.Errorf("load destination queue session: %w", err) - } - if to == nil { - return fmt.Errorf("load destination queue session %q: not found", toSessionID) - } - fromIdentity := messagequeue.QueueSessionIdentity{ - TaskID: from.TaskID, SessionID: from.ID, SessionIncarnationID: from.QueueIncarnationID, - } - if fromIdentity.SessionIncarnationID == "" { - fromIdentity, err = s.messageQueue.ResolveSessionIdentity(ctx, from.TaskID, from.ID) - if err != nil { - return fmt.Errorf("resolve source queue session: %w", err) - } - } - toIdentity := messagequeue.QueueSessionIdentity{ - TaskID: to.TaskID, SessionID: to.ID, SessionIncarnationID: to.QueueIncarnationID, - } - if toIdentity.SessionIncarnationID == "" { - toIdentity, err = s.messageQueue.ResolveSessionIdentity(ctx, to.TaskID, to.ID) - if err != nil { - return fmt.Errorf("resolve destination queue session: %w", err) - } - } - return s.messageQueue.TransferSessionIdentities(ctx, fromIdentity, toIdentity) -} - // completeAndStopSession stops the agent for a session and marks it COMPLETED. // Used by both the reuse path and the create-new path to terminate the // previous session in a uniform way. The state transition is the only session @@ -4253,6 +4458,9 @@ func (s *Service) dispatchTakenQueuedMessageForSession( s.logger.Warn("skipping empty queued message after transition", zap.String("session_id", identity.SessionID), zap.String("queue_id", queuedMsg.ID)) + if queuedMsg.IsDurableLifecycle() { + s.acknowledgeLifecycleQueueEntry(ctx, identity.SessionID, queuedMsg) + } return false } // Reserve entryID before handing off to the async goroutine. The worker @@ -4268,6 +4476,38 @@ func (s *Service) dispatchTakenQueuedMessageForSession( return true } +// dispatchTakenQueuedMessage keeps the source-session reservation path used by +// older queue callers and focused tests. New lifecycle code should pass the +// captured session identity to dispatchTakenQueuedMessageForSession. +func (s *Service) dispatchTakenQueuedMessage( + ctx context.Context, + sessionID string, + queuedMsg *messagequeue.QueuedMessage, + ok bool, +) bool { + if !ok || queuedMsg == nil { + return false + } + if queuedMsg.TaskID != "" { + if identity, err := s.messageQueue.ResolveSessionIdentity(ctx, queuedMsg.TaskID, sessionID); err == nil { + return s.dispatchTakenQueuedMessageForSession(ctx, identity, queuedMsg, ok) + } + } + s.publishQueueStatusEvent(ctx, sessionID) + if queuedMsg.Content == "" && len(queuedMsg.Attachments) == 0 { + s.logger.Warn("discarding empty queued message after transition", + zap.String("session_id", sessionID), + zap.String("queue_id", queuedMsg.ID)) + if queuedMsg.IsDurableLifecycle() { + s.acknowledgeLifecycleQueueEntry(ctx, sessionID, queuedMsg) + } + return false + } + reservation := s.markQueuedDispatchInFlightWithSourceLocked(sessionID, queuedMsg.ID, queuedMsg) + go s.executeQueuedMessageWithReservation(sessionID, queuedMsg, reservation) + return true +} + type passthroughRunningPreparer interface { PreparePassthroughRunning(sessionID string) (func(), error) } @@ -4897,7 +5137,13 @@ func (s *Service) takeAndMergeHandoffMessage(ctx context.Context, sessionID, bas if msg != nil { queuedHandoff = stepHandoffFromQueuedMetadata(msg.Metadata) } - if !ok || msg == nil || (msg.Content == "" && len(msg.Attachments) == 0 && queuedHandoff == "") { + if !ok || msg == nil { + return nil, basePrompt, nil, nil, "" + } + if msg.Content == "" && len(msg.Attachments) == 0 && queuedHandoff == "" { + s.logger.Warn("discarding empty hand-off queue entry", + zap.String("session_id", sessionID), + zap.String("queue_id", msg.ID)) return nil, basePrompt, nil, nil, "" } prompt := basePrompt diff --git a/apps/backend/internal/orchestrator/event_handlers_workflow_attachment_test.go b/apps/backend/internal/orchestrator/event_handlers_workflow_attachment_test.go new file mode 100644 index 0000000000..5d60b5a818 --- /dev/null +++ b/apps/backend/internal/orchestrator/event_handlers_workflow_attachment_test.go @@ -0,0 +1,729 @@ +package orchestrator + +import ( + "context" + "database/sql" + "errors" + "path/filepath" + "testing" + "time" + + "github.com/jmoiron/sqlx" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + taskservice "github.com/kandev/kandev/internal/task/service" + _ "github.com/mattn/go-sqlite3" +) + +type workflowAttachmentTransferStub struct { + calls []struct { + taskID, oldSessionID, newSessionID string + } + started chan struct{} + release chan struct{} +} + +func (s *workflowAttachmentTransferStub) TransferSessionMessageAttachments( + _ context.Context, + taskID, oldSessionID, newSessionID string, + _ []string, +) error { + s.calls = append(s.calls, struct { + taskID, oldSessionID, newSessionID string + }{taskID, oldSessionID, newSessionID}) + if s.started != nil { + close(s.started) + <-s.release + } + return nil +} + +func TestTransferQueuedSessionStateRebindsAttachments(t *testing.T) { + ctx := context.Background() + transfer := &workflowAttachmentTransferStub{} + svc := &Service{ + logger: testLogger(), + messageQueue: messagequeue.NewServiceMemory(testLogger()), + sessionAttachmentTransferer: transfer, + } + queued, err := svc.messageQueue.QueueMessage(ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}) + if err != nil { + t.Fatal(err) + } + + if err := svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new"); err != nil { + t.Fatal(err) + } + if len(transfer.calls) != 1 || transfer.calls[0].taskID != "task-transfer" || transfer.calls[0].oldSessionID != "session-old" || transfer.calls[0].newSessionID != "session-new" { + t.Fatalf("attachment transfer calls = %+v", transfer.calls) + } + moved, ok := svc.messageQueue.TakeQueued(ctx, "session-new") + if !ok || moved.ID != queued.ID { + t.Fatalf("moved queue entry = %#v, ok=%t", moved, ok) + } +} +func TestTransferQueuedSessionStateFailsClosedWithoutAttachmentTransferer(t *testing.T) { + ctx := context.Background() + queue := messagequeue.NewServiceMemory(testLogger()) + svc := &Service{logger: testLogger(), messageQueue: queue} + if _, err := queue.QueueMessage( + ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + ); err != nil { + t.Fatal(err) + } + + err := svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new") + if !errors.Is(err, errSessionAttachmentTransferUnavailable) { + t.Fatalf("transfer error = %v, want attachment transfer service unavailable", err) + } + if _, ok := queue.TakeQueued(ctx, "session-new"); ok { + t.Fatal("queue entry moved despite unavailable attachment transfer service") + } + if _, ok := queue.TakeQueued(ctx, "session-old"); !ok { + t.Fatal("queue entry missing after failed transfer") + } +} +func TestTransferQueuedSessionStateFailsClosedWhenTaskAttachmentServiceUnavailable(t *testing.T) { + ctx := context.Background() + queue := messagequeue.NewServiceMemory(testLogger()) + svc := &Service{ + logger: testLogger(), messageQueue: queue, + sessionAttachmentTransferer: &taskservice.Service{}, + } + if _, err := queue.QueueMessage( + ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + ); err != nil { + t.Fatal(err) + } + + if err := svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new"); err == nil { + t.Fatal("transfer unexpectedly succeeded without task attachment service") + } + if _, ok := queue.TakeQueued(ctx, "session-new"); ok { + t.Fatal("queue entry moved despite unavailable task attachment service") + } + if _, ok := queue.TakeQueued(ctx, "session-old"); !ok { + t.Fatal("queue entry missing after failed transfer") + } +} + +func TestTransferQueuedSessionStateTreatsTypedNilAttachmentTransfererAsUnavailable(t *testing.T) { + ctx := context.Background() + queue := messagequeue.NewServiceMemory(testLogger()) + var transfer *taskservice.Service + svc := &Service{ + logger: testLogger(), messageQueue: queue, + sessionAttachmentTransferer: transfer, + } + if _, err := queue.QueueMessage( + ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + ); err != nil { + t.Fatal(err) + } + + err := svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new") + if !errors.Is(err, errSessionAttachmentTransferUnavailable) { + t.Fatalf("transfer error = %v, want attachment transfer service unavailable", err) + } + if _, ok := queue.TakeQueued(ctx, "session-new"); ok { + t.Fatal("queue entry moved despite typed nil attachment transfer service") + } +} + +func TestSessionTransferRecoveryFailsClosedWhenTaskAttachmentServiceUnavailable(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + entry, err := queue.QueueMessage( + ctx, "session-new", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + if err := queue.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "unavailable-attachment-service", + TaskID: entry.TaskID, + FromSessionID: "session-old", + ToSessionID: entry.SessionID, + EntryIDs: []string{entry.ID}, + AttachmentIDs: []string{"attachment"}, + }); err != nil { + t.Fatal(err) + } + expireSessionTransferCompensationLease(t, db) + svc := &Service{ + logger: testLogger(), messageQueue: queue, + sessionAttachmentTransferer: &taskservice.Service{}, + } + + if err := svc.reconcileSessionTransferCompensationsOnStartup(ctx); err == nil { + t.Fatal("recovery unexpectedly succeeded without task attachment service") + } + compensations, err := queue.ListSessionTransferCompensations(ctx) + if err != nil { + t.Fatal(err) + } + if len(compensations) != 1 { + t.Fatalf("remaining compensations = %#v, want one", compensations) + } +} +func TestSessionTransferCompensationRecoveryAllowsTextOnlyWithoutAttachmentTransferer(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + entry, err := queue.QueueMessage( + ctx, "session-new", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + if err := queue.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "text-only-transfer", + TaskID: entry.TaskID, + FromSessionID: "session-old", + ToSessionID: entry.SessionID, + EntryIDs: []string{entry.ID}, + }); err != nil { + t.Fatal(err) + } + expireSessionTransferCompensationLease(t, db) + svc := &Service{logger: testLogger(), messageQueue: queue} + + if err := svc.reconcileSessionTransferCompensationsOnStartup(ctx); err != nil { + t.Fatal(err) + } + compensations, err := queue.ListSessionTransferCompensations(ctx) + if err != nil { + t.Fatal(err) + } + if len(compensations) != 0 { + t.Fatalf("remaining compensations = %#v, want none", compensations) + } +} + +func TestTransferQueuedSessionStateSerializesAttachmentTransferWithQueueMutation(t *testing.T) { + ctx := context.Background() + transfer := &workflowAttachmentTransferStub{ + started: make(chan struct{}), + release: make(chan struct{}), + } + svc := &Service{ + logger: testLogger(), + messageQueue: messagequeue.NewServiceMemory(testLogger()), + sessionAttachmentTransferer: transfer, + } + queued, err := svc.messageQueue.QueueMessage(ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}) + if err != nil { + t.Fatal(err) + } + + transferDone := make(chan error, 1) + go func() { + transferDone <- svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new") + }() + select { + case <-transfer.started: + case <-time.After(time.Second): + t.Fatal("attachment transfer did not start") + } + + mutationDone := make(chan error, 1) + go func() { + mutationDone <- svc.messageQueue.UpdateMessageWithMetadata( + ctx, "session-old", queued.ID, "changed", nil, nil, messagequeue.QueuedByUser, + ) + }() + select { + case err := <-mutationDone: + t.Fatalf("queue mutation completed during attachment transfer: %v", err) + case <-time.After(50 * time.Millisecond): + } + close(transfer.release) + + if err := <-transferDone; err != nil { + t.Fatal(err) + } + if err := <-mutationDone; err == nil { + t.Fatal("queue mutation unexpectedly succeeded after transfer") + } +} + +func TestTransferQueuedSessionStateIncludesReservedAttachmentRows(t *testing.T) { + ctx := context.Background() + transfer := &workflowAttachmentTransferStub{} + queue := messagequeue.NewServiceMemory(testLogger()) + svc := &Service{ + logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer, + } + queued, err := queue.QueueMessageWithMetadata( + ctx, + "session-old", + "task-transfer", + "reserved handoff", + "", + messagequeue.QueuedByAgent, + false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + map[string]interface{}{messagequeue.MetadataLifecycleDurable: true}, + ) + if err != nil { + t.Fatal(err) + } + if _, err := queue.ClaimSendNow(ctx, "session-old", []messagequeue.QueuedMessage{*queued}); err != nil { + t.Fatal(err) + } + if queue.GetStatus(ctx, "session-old").Count != 0 { + t.Fatal("reserved lifecycle row remained visible") + } + + if err := svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new"); err != nil { + t.Fatal(err) + } + if len(transfer.calls) != 1 { + t.Fatalf("attachment transfer calls = %d, want 1", len(transfer.calls)) + } +} + +type statefulWorkflowAttachmentTransfer struct { + currentSession string + rollbackErr error + calls int +} + +func (s *statefulWorkflowAttachmentTransfer) TransferSessionMessageAttachments( + _ context.Context, + _ string, + oldSessionID, newSessionID string, + _ []string, +) error { + s.calls++ + if s.currentSession != oldSessionID { + return errors.New("attachment session mismatch") + } + if s.calls > 1 && s.rollbackErr != nil { + return s.rollbackErr + } + s.currentSession = newSessionID + return nil +} + +func TestTransferQueuedSessionStateRecoversFailedAttachmentCompensationAfterRestart(t *testing.T) { + ctx := context.Background() + dbPath := filepath.Join(t.TempDir(), "queue.db") + queue, db := newWorkflowTransferQueue(t, dbPath) + _, err := queue.QueueMessage(ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}) + if err != nil { + t.Fatal(err) + } + if _, err := db.Exec(` + CREATE TRIGGER fail_workflow_queue_transfer + BEFORE UPDATE OF session_id ON queued_messages + BEGIN + SELECT RAISE(ABORT, 'forced queue transfer failure'); + END + `); err != nil { + t.Fatal(err) + } + transfer := &statefulWorkflowAttachmentTransfer{ + currentSession: "session-old", + rollbackErr: errors.New("forced attachment rollback failure"), + } + svc := &Service{logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer} + + if err := svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new"); err == nil { + t.Fatal("transfer unexpectedly succeeded") + } + if transfer.currentSession != "session-new" { + t.Fatalf("attachment session after failed compensation = %q, want session-new before recovery", transfer.currentSession) + } + expireSessionTransferCompensationLease(t, db) + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + transfer.rollbackErr = nil + restarted := &Service{ + logger: testLogger(), messageQueue: restartedQueue, sessionAttachmentTransferer: transfer, + } + reconciler, ok := any(restarted).(interface { + reconcileSessionTransferCompensationsOnStartup(context.Context) error + }) + if !ok { + t.Fatal("orchestrator has no durable session-transfer compensation reconciler") + } + if err := reconciler.reconcileSessionTransferCompensationsOnStartup(ctx); err != nil { + t.Fatal(err) + } + if transfer.currentSession != "session-old" { + t.Fatalf("attachment session after recovery = %q, want session-old", transfer.currentSession) + } +} + +type attachmentSessionSetTransfer struct { + sessions map[string]string +} + +func (s *attachmentSessionSetTransfer) TransferSessionMessageAttachments( + _ context.Context, + _ string, + oldSessionID, newSessionID string, + attachmentIDs []string, +) error { + for _, attachmentID := range attachmentIDs { + if s.sessions[attachmentID] == oldSessionID { + s.sessions[attachmentID] = newSessionID + } + } + return nil +} + +func TestTransferQueuedSessionRollbackPreservesDestinationAttachmentClaims(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + _, err := queue.QueueMessage( + ctx, "session-old", "task-transfer", "source", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "source-attachment"}}, + ) + if err != nil { + t.Fatal(err) + } + _, err = queue.QueueMessage( + ctx, "session-new", "task-transfer", "destination", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "destination-attachment"}}, + ) + if err != nil { + t.Fatal(err) + } + if _, err := db.Exec(` + CREATE TRIGGER fail_scoped_workflow_queue_transfer + BEFORE UPDATE OF session_id ON queued_messages + BEGIN + SELECT RAISE(ABORT, 'forced queue transfer failure'); + END + `); err != nil { + t.Fatal(err) + } + transfer := &attachmentSessionSetTransfer{sessions: map[string]string{ + "source-attachment": "session-old", + "destination-attachment": "session-new", + }} + svc := &Service{logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer} + + if err := svc.transferQueuedSessionState(ctx, "task-transfer", "session-old", "session-new"); err == nil { + t.Fatal("transfer unexpectedly succeeded") + } + if got := transfer.sessions["destination-attachment"]; got != "session-new" { + t.Fatalf("pre-existing destination attachment session = %q, want session-new", got) + } +} +func TestSessionTransferCompensationRecoveryAcceptsDestinationAttachment(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + entry, err := queue.QueueMessage( + ctx, "session-new", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + if err := queue.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "destination-transfer", + TaskID: entry.TaskID, + FromSessionID: "session-old", + ToSessionID: entry.SessionID, + EntryIDs: []string{entry.ID}, + AttachmentIDs: []string{"destination-attachment"}, + }); err != nil { + t.Fatal(err) + } + expireSessionTransferCompensationLease(t, db) + transfer := &attachmentSessionSetTransfer{sessions: map[string]string{ + "destination-attachment": "session-new", + }} + restarted := &Service{ + logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer, + } + if err := restarted.reconcileSessionTransferCompensationsOnStartup(ctx); err != nil { + t.Fatal(err) + } + if got := transfer.sessions["destination-attachment"]; got != "session-new" { + t.Fatalf("destination attachment session = %q, want session-new", got) + } + compensations, err := queue.ListSessionTransferCompensations(ctx) + if err != nil { + t.Fatal(err) + } + if len(compensations) != 0 { + t.Fatalf("remaining compensations = %#v, want none", compensations) + } +} + +func TestCleanupOnlyTransferCompensationRecoversSourceAfterRestart(t *testing.T) { + ctx := context.Background() + dbPath := filepath.Join(t.TempDir(), "queue.db") + queue, db := newWorkflowTransferQueue(t, dbPath) + cleanup := messagequeue.AttachmentCleanup{ + SessionID: "session-old", EntryID: "cleanup-entry", OperationID: "cleanup-operation", + TaskID: "task-transfer", Attachments: []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + } + if err := queue.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + if _, err := db.Exec(` + CREATE TRIGGER fail_cleanup_only_queue_transfer + BEFORE UPDATE OF current_session_id ON queue_attachment_cleanups + BEGIN + SELECT RAISE(ABORT, 'forced cleanup transfer failure'); + END + `); err != nil { + t.Fatal(err) + } + transfer := &statefulWorkflowAttachmentTransfer{ + currentSession: "session-old", + rollbackErr: errors.New("forced attachment rollback failure"), + } + svc := &Service{logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer} + if err := svc.transferQueuedSessionState(ctx, cleanup.TaskID, cleanup.SessionID, "session-new"); err == nil { + t.Fatal("transfer unexpectedly succeeded") + } + if transfer.currentSession != "session-new" { + t.Fatalf("attachment session before restart = %q, want session-new", transfer.currentSession) + } + expireSessionTransferCompensationLease(t, db) + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + transfer.rollbackErr = nil + restarted := &Service{ + logger: testLogger(), messageQueue: restartedQueue, sessionAttachmentTransferer: transfer, + } + if err := restarted.reconcileSessionTransferCompensationsOnStartup(ctx); err != nil { + t.Fatal(err) + } + if transfer.currentSession != "session-old" { + t.Fatalf("attachment session after recovery = %q, want session-old", transfer.currentSession) + } +} + +func TestSessionTransferCompensationRecoveryRejectsActiveOwner(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + entry, err := queue.QueueMessage( + ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + ) + if err != nil { + t.Fatal(err) + } + if err := queue.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "active-transfer", + TaskID: entry.TaskID, + FromSessionID: entry.SessionID, + ToSessionID: "session-new", + EntryIDs: []string{entry.ID}, + AttachmentIDs: []string{"attachment"}, + }); err != nil { + t.Fatal(err) + } + transfer := &workflowAttachmentTransferStub{} + restarted := &Service{ + logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer, + } + + err = restarted.reconcileSessionTransferCompensationsOnStartup(ctx) + + if !errors.Is(err, messagequeue.ErrSessionTransferInProgress) { + t.Fatalf("recovery error = %v, want active transfer rejection", err) + } + if len(transfer.calls) != 0 { + t.Fatalf("attachment recovery calls = %#v, want none", transfer.calls) + } + compensations, listErr := queue.ListSessionTransferCompensations(ctx) + if listErr != nil { + t.Fatal(listErr) + } + if len(compensations) != 1 { + t.Fatalf("remaining compensations = %#v, want active owner preserved", compensations) + } +} + +func TestDurableQueueStartupRecoveryStopsRetryWhenContextCancels(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + queue, db := newWorkflowTransferQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + entry, err := queue.QueueMessage( + ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + if err := queue.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "active-transfer", TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + service := &Service{ + logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: &workflowAttachmentTransferStub{}, + } + result := make(chan error, 1) + go func() { result <- service.reconcileDurableQueueStateOnStartup(ctx) }() + + select { + case err := <-result: + t.Fatalf("startup recovery returned before cancellation: %v", err) + case <-time.After(20 * time.Millisecond): + } + cancel() + select { + case err := <-result: + if !errors.Is(err, context.Canceled) { + t.Fatalf("startup recovery error = %v, want context cancellation", err) + } + case <-time.After(time.Second): + t.Fatal("startup recovery did not stop after cancellation") + } +} + +func TestCleanupOnlyTransferCompensationRecoversPreviousHopAfterRestart(t *testing.T) { + ctx := context.Background() + dbPath := filepath.Join(t.TempDir(), "queue.db") + queue, db := newWorkflowTransferQueue(t, dbPath) + cleanup := messagequeue.AttachmentCleanup{ + SessionID: "session-a", EntryID: "cleanup-entry", OperationID: "cleanup-operation", + TaskID: "task-transfer", Attachments: []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + } + if err := queue.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + transfer := &statefulWorkflowAttachmentTransfer{currentSession: "session-a"} + svc := &Service{logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer} + if err := svc.transferQueuedSessionState(ctx, cleanup.TaskID, "session-a", "session-b"); err != nil { + t.Fatal(err) + } + if _, err := db.Exec(` + CREATE TRIGGER fail_multi_hop_cleanup_transfer + BEFORE UPDATE OF current_session_id ON queue_attachment_cleanups + BEGIN + SELECT RAISE(ABORT, 'forced cleanup transfer failure'); + END + `); err != nil { + t.Fatal(err) + } + transfer.calls = 0 + transfer.rollbackErr = errors.New("forced attachment rollback failure") + if err := svc.transferQueuedSessionState(ctx, cleanup.TaskID, "session-b", "session-c"); err == nil { + t.Fatal("second transfer unexpectedly succeeded") + } + if transfer.currentSession != "session-c" { + t.Fatalf("attachment session before restart = %q, want session-c", transfer.currentSession) + } + expireSessionTransferCompensationLease(t, db) + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + transfer.rollbackErr = nil + restarted := &Service{ + logger: testLogger(), messageQueue: restartedQueue, sessionAttachmentTransferer: transfer, + } + if err := restarted.reconcileSessionTransferCompensationsOnStartup(ctx); err != nil { + t.Fatal(err) + } + if transfer.currentSession != "session-b" { + t.Fatalf("attachment session after recovery = %q, want session-b", transfer.currentSession) + } +} + +func TestDurableQueueStartupRecoveryReconcilesTransferBeforeDispatch(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + entry, err := queue.QueueMessage( + ctx, "session-old", "task-transfer", "handoff", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := queue.ReserveQueued(ctx, entry.SessionID) + if !ok || reserved.ID != entry.ID { + t.Fatalf("reserved entry = %#v, ok=%t", reserved, ok) + } + if err := queue.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "startup-transfer", + TaskID: entry.TaskID, + FromSessionID: entry.SessionID, + ToSessionID: "session-new", + EntryIDs: []string{entry.ID}, + AttachmentIDs: []string{"attachment"}, + CleanupLocators: nil, + }); err != nil { + t.Fatal(err) + } + expireSessionTransferCompensationLease(t, db) + transfer := &statefulWorkflowAttachmentTransfer{currentSession: "session-new"} + restarted := &Service{ + logger: testLogger(), messageQueue: queue, sessionAttachmentTransferer: transfer, + } + + if err := restarted.reconcileDurableQueueStateOnStartup(ctx); err != nil { + t.Fatal(err) + } + + if transfer.currentSession != entry.SessionID { + t.Fatalf("attachment session after recovery = %q, want %q", transfer.currentSession, entry.SessionID) + } + status := queue.GetStatus(ctx, entry.SessionID) + if len(status.Entries) != 1 || status.Entries[0].ID != entry.ID { + t.Fatalf("recovered source queue = %#v", status.Entries) + } + compensations, err := queue.ListSessionTransferCompensations(ctx) + if err != nil { + t.Fatal(err) + } + if len(compensations) != 0 { + t.Fatalf("remaining compensations = %#v", compensations) + } +} + +func expireSessionTransferCompensationLease(t *testing.T, db *sqlx.DB) { + t.Helper() + if _, err := db.Exec(` + UPDATE queue_session_transfer_compensations + SET recovery_lease_expires_at = ? + `, time.Now().UTC().Add(-time.Second)); err != nil { + t.Fatal(err) + } +} + +func newWorkflowTransferQueue(t *testing.T, dbPath string) (*messagequeue.Service, *sqlx.DB) { + t.Helper() + raw, err := sql.Open("sqlite3", dbPath) + if err != nil { + t.Fatal(err) + } + + raw.SetMaxOpenConns(1) + db := sqlx.NewDb(raw, "sqlite3") + repo, err := messagequeue.NewSQLiteRepository(db, db) + if err != nil { + _ = db.Close() + t.Fatal(err) + } + return messagequeue.NewService(repo, messagequeue.DefaultMaxPerSession, testLogger()), db +} diff --git a/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go b/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go index b579c4b499..527062bd59 100644 --- a/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go @@ -117,6 +117,89 @@ func TestCreateNewSessionForStepKeepsCurrentSessionWhenWorkspaceAttachFails(t *t } } +type transferFailureQueueRepository struct { + messagequeue.Repository + err error +} + +func (r *transferFailureQueueRepository) TransferSession(context.Context, string, string) error { + return r.err +} + +func (r *transferFailureQueueRepository) TransferSessionIdentities( + context.Context, + messagequeue.QueueSessionIdentity, + messagequeue.QueueSessionIdentity, +) error { + return r.err +} + +func TestCreateNewSessionForStepFailsClosedWhenQueueTransferFails(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-queue-transfer", "session-queue-transfer", "step-one") + current, err := repo.GetTaskSession(ctx, "session-queue-transfer") + if err != nil { + t.Fatal(err) + } + current.State = models.TaskSessionStateRunning + current.IsPrimary = true + current.AgentProfileID = "profile-old" + current.ExecutorID = models.ExecutorIDWorktree + current.TaskEnvironmentID = "environment-queue-transfer" + if err := repo.UpdateTaskSession(ctx, current); err != nil { + t.Fatal(err) + } + if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{ + ID: current.TaskEnvironmentID, TaskID: current.TaskID, + ExecutorType: string(models.ExecutorTypeLocal), Status: models.TaskEnvironmentStatusReady, + }); err != nil { + t.Fatal(err) + } + taskRepo := newMockTaskRepo() + taskRepo.tasks[current.TaskID] = &v1.Task{ID: current.TaskID, WorkspaceID: "ws1", Title: "Test Task"} + agentMgr := &mockAgentManager{ + repoForExecutionLookup: repo, + launchAgentFunc: func(context.Context, *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) { + return &executor.LaunchAgentResponse{AgentExecutionID: "replacement-execution"}, nil + }, + } + svc := createTestServiceWithScheduler(repo, newMockStepGetter(), taskRepo, agentMgr) + transferErr := errors.New("queue transfer failed") + svc.messageQueue = messagequeue.NewService( + &transferFailureQueueRepository{Repository: newAuthoritativeMemoryRepository(repo), err: transferErr}, + messagequeue.DefaultMaxPerSession, + testLogger(), + ) + if _, err := svc.messageQueue.QueueMessage(ctx, current.ID, current.TaskID, "handoff", "", messagequeue.QueuedByUser, false, nil); err != nil { + t.Fatal(err) + } + + _, err = svc.createNewSessionForStep(ctx, current.TaskID, current, "profile-new") + if !errors.Is(err, transferErr) { + t.Fatalf("createNewSessionForStep error = %v, want queue transfer failure", err) + } + persisted, err := repo.GetTaskSession(ctx, current.ID) + if err != nil { + t.Fatal(err) + } + if persisted.State != models.TaskSessionStateRunning || !persisted.IsPrimary { + t.Fatalf("current session after failed queue transfer = state %q primary %t, want running primary", persisted.State, persisted.IsPrimary) + } + if got := svc.messageQueue.GetStatus(ctx, current.ID).Count; got != 1 { + t.Fatalf("queued hand-off after failed transfer = %d, want 1", got) + } + sessions, err := repo.ListTaskSessions(ctx, current.TaskID) + if err != nil { + t.Fatal(err) + } + for _, session := range sessions { + if session.ID != current.ID && session.State != models.TaskSessionStateCompleted { + t.Fatalf("replacement session after failed transfer = %s, want completed cleanup", session.State) + } + } +} + // terminalizeCandidateBeforePromotionRepo pauses a profile-switch promotion // after lookup so the test can terminalize the selected row before the // promotion write begins. diff --git a/apps/backend/internal/orchestrator/event_handlers_workflow_replacement_queue_test.go b/apps/backend/internal/orchestrator/event_handlers_workflow_replacement_queue_test.go new file mode 100644 index 0000000000..746fcf30ba --- /dev/null +++ b/apps/backend/internal/orchestrator/event_handlers_workflow_replacement_queue_test.go @@ -0,0 +1,95 @@ +package orchestrator + +import ( + "context" + "database/sql" + "path/filepath" + "testing" + + "github.com/jmoiron/sqlx" + _ "github.com/mattn/go-sqlite3" + + v1 "github.com/kandev/kandev/pkg/api/v1" + + "github.com/kandev/kandev/internal/orchestrator/executor" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/task/models" +) + +func TestCreateNewSessionForStepTransfersQueueStateExactlyOnce(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-queue-state", "session-queue-state", "step-one") + current, err := repo.GetTaskSession(ctx, "session-queue-state") + if err != nil { + t.Fatal(err) + } + current.State = models.TaskSessionStateRunning + current.IsPrimary = true + current.AgentProfileID = "profile-old" + current.ExecutorID = models.ExecutorIDWorktree + current.TaskEnvironmentID = "environment-queue-state" + if err := repo.UpdateTaskSession(ctx, current); err != nil { + t.Fatal(err) + } + if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{ + ID: current.TaskEnvironmentID, TaskID: current.TaskID, + ExecutorType: string(models.ExecutorTypeLocal), Status: models.TaskEnvironmentStatusReady, + }); err != nil { + t.Fatal(err) + } + taskRepo := newMockTaskRepo() + taskRepo.tasks[current.TaskID] = &v1.Task{ID: current.TaskID, WorkspaceID: "ws1", Title: "Test Task"} + svc := createTestServiceWithScheduler(repo, newMockStepGetter(), taskRepo, &mockAgentManager{ + repoForExecutionLookup: repo, + launchAgentFunc: func(context.Context, *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) { + return &executor.LaunchAgentResponse{AgentExecutionID: "replacement-execution"}, nil + }, + }) + + rawDB, err := sql.Open("sqlite3", filepath.Join(t.TempDir(), "queue.db")) + if err != nil { + t.Fatal(err) + } + db := sqlx.NewDb(rawDB, "sqlite3") + t.Cleanup(func() { _ = db.Close() }) + queueRepo, err := messagequeue.NewSQLiteRepository(db, db) + if err != nil { + t.Fatal(err) + } + svc.messageQueue = messagequeue.NewService(queueRepo, messagequeue.DefaultMaxPerSession, testLogger()) + attachmentTransfer := &workflowAttachmentTransferStub{} + svc.SetSessionAttachmentTransferer(attachmentTransfer) + queued, err := svc.messageQueue.QueueMessage(ctx, current.ID, current.TaskID, "attachment-backed", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment-one", Name: "one.txt", MimeType: "text/plain"}}) + if err != nil { + t.Fatal(err) + } + if err := svc.messageQueue.SetPendingMove(ctx, current.ID, &messagequeue.PendingMove{ + MoveID: "move-one", TaskID: current.TaskID, WorkflowID: "workflow-one", WorkflowStepID: "step-two", + }); err != nil { + t.Fatal(err) + } + + created, err := svc.createNewSessionForStep(ctx, current.TaskID, current, "profile-new") + if err != nil { + t.Fatalf("createNewSessionForStep: %v", err) + } + move, _, err := svc.messageQueue.GetPendingMoveWithError(ctx, created.ID) + if err != nil { + t.Fatal(err) + } + if move == nil || move.MoveID != "move-one" { + t.Fatalf("replacement pending move = %+v, want move-one", move) + } + transferred, err := svc.messageQueue.GetEntry(ctx, created.ID, queued.ID) + if err != nil { + t.Fatal(err) + } + if len(transferred.Attachments) != 1 || transferred.Attachments[0].AttachmentID != "attachment-one" { + t.Fatalf("replacement attachments = %+v, want attachment-one", transferred.Attachments) + } + if len(attachmentTransfer.calls) != 1 { + t.Fatalf("attachment transfer calls = %d, want 1", len(attachmentTransfer.calls)) + } +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers.go b/apps/backend/internal/orchestrator/handlers/queue_handlers.go index 8a6bd87844..df125395cd 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_handlers.go +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers.go @@ -2,10 +2,14 @@ package handlers import ( "context" + "crypto/sha256" "encoding/base64" + "encoding/hex" + "encoding/json" "errors" "fmt" - + "github.com/google/uuid" + "github.com/kandev/kandev/internal/auth/authn" "github.com/kandev/kandev/internal/common/logger" "github.com/kandev/kandev/internal/entityrefs" "github.com/kandev/kandev/internal/events/bus" @@ -15,6 +19,8 @@ import ( v1 "github.com/kandev/kandev/pkg/api/v1" ws "github.com/kandev/kandev/pkg/websocket" "go.uber.org/zap" + "sync" + "time" ) const ( @@ -154,6 +160,58 @@ type QueueAutoMergeController interface { SetSessionAutoMerge(context.Context, messagequeue.QueueSessionIdentity, bool) (messagequeue.AutoMergePolicy, error) } +// QueueEditLeaseController owns target-bound queued-message edit leases. +type QueueEditLeaseController interface { + BeginEdit(context.Context, string, string, string) (*messagequeue.QueueEditLease, error) + RenewEdit(context.Context, string, string, string, string) (*messagequeue.QueueEditLease, error) + EndEdit(context.Context, string, string, string, string) error + UpdateMessageWithLease(context.Context, string, string, string, string, string, int64, string, []messagequeue.MessageAttachment, map[string]interface{}) (int64, error) +} +type queueEditSavedLeaseController interface { + EndEditAfterSave(context.Context, string, string, string, string) (bool, error) +} + +// QueueAutoRunPreservingDrainer attempts one normal FIFO dispatch without +// changing the persisted Auto-run policy. +type QueueAutoRunPreservingDrainer interface { + DrainQueuedMessageIfAutoRun(context.Context, string) (bool, error) +} + +type queueEditAdmissionController interface { + WithSessionAdmission(context.Context, string, func(context.Context) error) error +} +type queueEditLeaseStateReader interface { + GetEditLease(context.Context, string, string) (*messagequeue.QueueEditLease, error) +} + +type queueEntryLocator interface { + FindEntryByID(context.Context, string) (*messagequeue.QueuedMessage, error) +} + +type queueAttachmentReferenceChecker interface { + ReferencedQueueAttachmentIDs(context.Context, string, string, []string) (map[string]struct{}, error) +} +type queueBatchCanceller interface { + CancelAllWithEntries(context.Context, string) ([]messagequeue.QueuedMessage, error) +} +type queueAttachmentCleanupStore interface { + AttachmentCleanupPersistenceAvailable() bool + UpsertAttachmentCleanup(context.Context, messagequeue.AttachmentCleanup) error + DeleteAttachmentCleanup(context.Context, string, string, string) error + ListAttachmentCleanups(context.Context) ([]messagequeue.AttachmentCleanup, error) +} +type queueAttachmentCleanupLocator interface { + GetAttachmentCleanup(context.Context, string, string, string) (*messagequeue.AttachmentCleanup, error) +} + +type queueEntryRemover interface { + RemoveEntryWithEntry(context.Context, string, string) (*messagequeue.QueuedMessage, error) +} + +type queueEditAttachmentController interface { + UpdateMessageWithLeaseAfterValidationAndFinalize(context.Context, string, string, string, string, string, int64, string, []messagequeue.MessageAttachment, map[string]interface{}, func(context.Context) error, func(context.Context) error, func(context.Context, *messagequeue.QueuedMessage) error) (int64, error) +} + // QueueSendNowDispatcher is implemented by the orchestrator service. It is // kept separate from QueueDrainer so queue-focused handlers can retain their // small test doubles while the new action gets the replacement-turn contract. @@ -193,27 +251,74 @@ type QueueAttachmentReleaser interface { ReleaseMessageAttachments(ctx context.Context, taskID, sessionID string, attachments []v1.MessageAttachment) error } +type queueAttachmentCleanupOwnerResolver interface { + ResolveMessageAttachmentOwner( + context.Context, string, string, []v1.MessageAttachment, + ) (string, error) +} + +type queueAttachmentCleanupReleaser interface { + ReleaseMessageAttachmentsForCleanup( + context.Context, string, string, []v1.MessageAttachment, + ) error +} + +type internalQueueAttachmentReleaser struct { + releaser queueAttachmentCleanupReleaser +} + +func (r internalQueueAttachmentReleaser) ReleaseMessageAttachments( + ctx context.Context, taskID, sessionID string, attachments []v1.MessageAttachment, +) error { + return r.releaser.ReleaseMessageAttachmentsForCleanup(ctx, taskID, sessionID, attachments) +} + type queueEntryTaker interface { TakeQueuedEntry(context.Context, string, string) (*messagequeue.QueuedMessage, bool, error) } +type pendingQueueAttachmentCleanupKey struct { + sessionID string + entryID string + operationID string +} + +type pendingQueueAttachmentCleanup struct { + key pendingQueueAttachmentCleanupKey + req wsUpdateMessageRequest + previous *messagequeue.QueuedMessage + releaser QueueAttachmentReleaser + removeEntry bool + claimPending bool + entryFingerprint string + currentSessionID string + authCtx context.Context + wake chan struct{} +} // QueueHandlers handles WebSocket message-queue operations. type QueueHandlers struct { - queueService QueueService - queueDrainer QueueDrainer - queueReadiness QueueAdmissionReadinessChecker - queueAutoRun QueueAutoRunController - queueDispatcher QueueSendNowDispatcher - accessAuthorizer QueueAccessAuthorizer - sessionTaskResolver SessionTaskResolver - eventBus bus.EventBus - logger *logger.Logger - referenceValidator entityrefs.SubmissionValidator - attachmentClaimer QueueAttachmentClaimer + queueService QueueService + queueDrainer QueueDrainer + queueReadiness QueueAdmissionReadinessChecker + queueAutoRun QueueAutoRunController + queueDispatcher QueueSendNowDispatcher + queueEdit QueueEditLeaseController + accessAuthorizer QueueAccessAuthorizer + sessionTaskResolver SessionTaskResolver + eventBus bus.EventBus + logger *logger.Logger + referenceValidator entityrefs.SubmissionValidator + attachmentClaimer QueueAttachmentClaimer + attachmentCleanupMu sync.Mutex + pendingAttachmentCleanup map[pendingQueueAttachmentCleanupKey]*pendingQueueAttachmentCleanup + attachmentCleanupCtx context.Context + attachmentCleanupCancel context.CancelFunc + attachmentCleanupWG sync.WaitGroup + attachmentCleanupStarted bool + attachmentCleanupStopped bool } -// SetAttachmentClaimer wires the task attachment registry into queue adds. -// It is optional so queue-focused tests and non-task consumers remain small. +// SetAttachmentClaimer wires task-owned attachment claiming into queue edits. func (h *QueueHandlers) SetAttachmentClaimer(claimer QueueAttachmentClaimer) { h.attachmentClaimer = claimer } @@ -235,13 +340,17 @@ func NewQueueHandlers( referenceValidator = validators[0] } handlers := &QueueHandlers{ - queueService: queueService, - queueDrainer: queueDrainer, - accessAuthorizer: accessAuthorizer, - sessionTaskResolver: sessionTaskResolver, - eventBus: eventBus, - logger: log.WithFields(zap.String("component", "queue-handlers")), - referenceValidator: referenceValidator, + queueService: queueService, + queueDrainer: queueDrainer, + accessAuthorizer: accessAuthorizer, + sessionTaskResolver: sessionTaskResolver, + eventBus: eventBus, + logger: log.WithFields(zap.String("component", "queue-handlers")), + referenceValidator: referenceValidator, + pendingAttachmentCleanup: make(map[pendingQueueAttachmentCleanupKey]*pendingQueueAttachmentCleanup), + } + if controller, ok := queueService.(QueueEditLeaseController); ok { + handlers.queueEdit = controller } if dispatcher, ok := queueDrainer.(QueueSendNowDispatcher); ok { handlers.queueDispatcher = dispatcher @@ -252,15 +361,142 @@ func NewQueueHandlers( if controller, ok := queueDrainer.(QueueAutoRunController); ok { handlers.queueAutoRun = controller } + return handlers } +// Start owns the context used by pending attachment cleanup retries and reloads +// durable obligations left by an earlier process. +func (h *QueueHandlers) Start(ctx context.Context) { + h.attachmentCleanupMu.Lock() + if h.attachmentCleanupStarted || h.attachmentCleanupStopped { + h.attachmentCleanupMu.Unlock() + return + } + h.attachmentCleanupCtx, h.attachmentCleanupCancel = context.WithCancel(ctx) + h.attachmentCleanupStarted = true + h.loadPendingAttachmentCleanupsLocked(context.WithoutCancel(ctx)) + for _, pending := range h.pendingAttachmentCleanup { + h.attachmentCleanupWG.Add(1) + go h.retryPendingAttachmentCleanup(pending) + } + h.attachmentCleanupMu.Unlock() +} + +func (h *QueueHandlers) loadPendingAttachmentCleanupsLocked(ctx context.Context) { + store, ok := h.attachmentCleanupStore() + releaser, canRelease := h.attachmentClaimer.(QueueAttachmentReleaser) + if !ok || !canRelease { + return + } + cleanups, err := store.ListAttachmentCleanups(ctx) + if err != nil { + h.logger.Error("failed to reload pending queue attachment cleanup", zap.Error(err)) + return + } + for _, cleanup := range cleanups { + key := pendingQueueAttachmentCleanupKey{ + sessionID: cleanup.SessionID, entryID: cleanup.EntryID, operationID: cleanup.OperationID, + } + if _, exists := h.pendingAttachmentCleanup[key]; exists { + continue + } + currentSessionID := cleanup.CurrentSessionID + if currentSessionID == "" { + currentSessionID = cleanup.SessionID + } + authCtx, cleanupReleaser := h.recoverAttachmentCleanupContext( + ctx, cleanup, currentSessionID, releaser, + ) + if cleanupReleaser == nil { + continue + } + h.pendingAttachmentCleanup[key] = &pendingQueueAttachmentCleanup{ + key: key, + req: wsUpdateMessageRequest{ + SessionID: cleanup.SessionID, + EntryID: cleanup.EntryID, + LeaseID: cleanup.LeaseID, + OperationID: cleanup.OperationID, + }, + previous: &messagequeue.QueuedMessage{ + ID: cleanup.EntryID, SessionID: cleanup.SessionID, TaskID: cleanup.TaskID, + Attachments: cleanup.Attachments, + }, + releaser: cleanupReleaser, + removeEntry: cleanup.RemoveEntry, + claimPending: cleanup.ClaimPending, + entryFingerprint: cleanup.EntryFingerprint, + currentSessionID: cleanup.CurrentSessionID, + authCtx: authCtx, + wake: make(chan struct{}, 1), + } + } +} + +func (h *QueueHandlers) recoverAttachmentCleanupContext( + ctx context.Context, + cleanup messagequeue.AttachmentCleanup, + currentSessionID string, + releaser QueueAttachmentReleaser, +) (context.Context, QueueAttachmentReleaser) { + authCtx := context.WithoutCancel(ctx) + var ownerID string + var resolveErr error + if cleanup.OwnerID != "" { + ownerID = cleanup.OwnerID + } else if resolver, resolverOK := h.attachmentClaimer.(queueAttachmentCleanupOwnerResolver); resolverOK { + ownerID, resolveErr = resolver.ResolveMessageAttachmentOwner( + authCtx, cleanup.TaskID, currentSessionID, queueAttachmentsToV1(cleanup.Attachments), + ) + } + var cleanupReleaser QueueAttachmentReleaser + if ownerID != "" { + authCtx = authn.WithIdentity(authCtx, authn.Identity{UserID: ownerID}) + cleanupReleaser = releaser + } else if cleanup.OwnerID == "" && !cleanup.ClaimPending { + if internalReleaser, canRelease := h.attachmentClaimer.(queueAttachmentCleanupReleaser); canRelease { + cleanupReleaser = internalQueueAttachmentReleaser{releaser: internalReleaser} + } + } + if cleanupReleaser == nil && cleanup.OwnerID == "" { + if resolveErr != nil { + h.logger.Warn("failed to recover owner for queue attachment cleanup", zap.Error(resolveErr)) + } else { + h.logger.Warn("queue attachment cleanup owner is unavailable") + } + } + return authCtx, cleanupReleaser +} + +func (h *QueueHandlers) attachmentCleanupStore() (queueAttachmentCleanupStore, bool) { + store, ok := h.queueService.(queueAttachmentCleanupStore) + return store, ok && store.AttachmentCleanupPersistenceAvailable() +} + +// Stop cancels active retries after their obligations are durable. A later +// handler instance reloads unresolved work before accepting queue operations. +func (h *QueueHandlers) Stop() { + h.attachmentCleanupMu.Lock() + if !h.attachmentCleanupStopped { + h.attachmentCleanupStopped = true + if h.attachmentCleanupCancel != nil { + h.attachmentCleanupCancel() + } + } + h.attachmentCleanupMu.Unlock() + h.attachmentCleanupWG.Wait() +} + // RegisterHandlers registers queue handlers with the dispatcher. func (h *QueueHandlers) RegisterHandlers(d *ws.Dispatcher) { d.RegisterFunc(ws.ActionMessageQueueAdd, h.wsQueueMessage) d.RegisterFunc(ws.ActionMessageQueueCancel, h.wsCancelAll) d.RegisterFunc(ws.ActionMessageQueueGet, h.wsGetQueueStatus) d.RegisterFunc(ws.ActionMessageQueueUpdate, h.wsUpdateMessage) + d.RegisterFunc(ws.ActionMessageQueueEditBegin, h.wsBeginEdit) + d.RegisterFunc(ws.ActionMessageQueueEditRenew, h.wsRenewEdit) + d.RegisterFunc(ws.ActionMessageQueueEditEnd, h.wsEndEdit) d.RegisterFunc(ws.ActionMessageQueueAppend, h.wsAppendToQueue) d.RegisterFunc(ws.ActionMessageQueueDrain, h.wsDrainQueue) d.RegisterFunc(ws.ActionMessageQueueSendNow, h.wsSendNow) @@ -370,8 +606,10 @@ func (h *QueueHandlers) wsQueueMessage(ctx context.Context, msg *ws.Message) (*w } var ( - errQueuedAttachmentUnavailable = errors.New("queued attachment unavailable") - errQueuedAttachmentRollback = errors.New("queued attachment rollback failed") + errQueuedAttachmentUnavailable = errors.New("queued attachment unavailable") + errQueuedAttachmentRollback = errors.New("queued attachment rollback failed") + errAttachmentCleanupLeaseActive = errors.New("queue attachment cleanup blocked by edit lease") + errAttachmentCleanupEntryChanged = errors.New("queue attachment cleanup entry changed") ) func (h *QueueHandlers) admitQueuedMessage(ctx context.Context, req *wsQueueMessageRequest, queuedBy string, metadata map[string]interface{}) (*messagequeue.QueuedMessage, error) { @@ -383,7 +621,63 @@ func (h *QueueHandlers) admitQueuedMessage(ctx context.Context, req *wsQueueMess } return h.queueService.QueueMessageWithMetadataAfterInsert( ctx, req.SessionID, req.TaskID, req.Content, req.Model, queuedBy, req.PlanMode, req.Attachments, metadata, - h.claimQueuedAttachmentsAfterInsert(req), + func(admittedCtx context.Context, source *messagequeue.QueuedMessage) error { + releaser, canRelease := h.attachmentClaimer.(QueueAttachmentReleaser) + if !canRelease { + return errors.New("queue attachment claimer cannot release failed claims") + } + pending, err := h.preparePendingAttachmentCleanupWithState( + admittedCtx, + wsUpdateMessageRequest{ + SessionID: source.SessionID, EntryID: source.ID, + OperationID: attachmentCleanupOperationID("", "admission"), + }, + source.TaskID, + source.Attachments, + releaser, + true, + source, + ) + if err != nil { + _, rollbackErr := h.rollbackQueuedAttachmentClaim(admittedCtx, source.SessionID, source.ID) + return fmt.Errorf("%w: %v", errQueuedAttachmentRollback, errors.Join(err, rollbackErr)) + } + claimErr := h.attachmentClaimer.ClaimMessageAttachments( + admittedCtx, req.TaskID, req.SessionID, queueAttachmentsToV1(req.Attachments), + ) + if claimErr == nil { + deleteErr := h.deletePendingAttachmentCleanup(pending) + if deleteErr == nil { + return nil + } + _, rollbackErr := h.rollbackQueuedAttachmentClaim(admittedCtx, req.SessionID, source.ID) + cleanupErr := h.releaseSupersededQueueAttachmentsAdmitted( + admittedCtx, pending.req, pending.previous, pending.releaser, + ) + h.settlePendingAttachmentCleanup(pending, errors.Join(rollbackErr, cleanupErr)) + return fmt.Errorf( + "%w: %v", + errQueuedAttachmentRollback, + errors.Join(deleteErr, rollbackErr, cleanupErr), + ) + } + pending.claimPending = false + pending.removeEntry = true + if err := h.persistPendingAttachmentCleanup(admittedCtx, pending); err != nil { + _, rollbackErr := h.rollbackQueuedAttachmentClaim(admittedCtx, req.SessionID, source.ID) + h.queuePendingAttachmentCleanup(pending) + return fmt.Errorf("%w: %v", errQueuedAttachmentRollback, errors.Join(err, rollbackErr)) + } + _, rollbackErr := h.rollbackQueuedAttachmentClaim(admittedCtx, req.SessionID, source.ID) + cleanupErr := h.releaseSupersededQueueAttachmentsAdmitted( + admittedCtx, pending.req, pending.previous, pending.releaser, + ) + h.settlePendingAttachmentCleanup(pending, errors.Join(rollbackErr, cleanupErr)) + if rollbackErr != nil || cleanupErr != nil { + return fmt.Errorf("%w: %v", errQueuedAttachmentRollback, errors.Join(rollbackErr, cleanupErr)) + } + return fmt.Errorf("%w: %v", errQueuedAttachmentUnavailable, claimErr) + }, ) } admissions, ok := h.queueService.(QueueIdentityAdmissionService) @@ -425,7 +719,7 @@ func (h *QueueHandlers) claimQueuedAttachmentsAfterInsert(req *wsQueueMessageReq if claimErr == nil { return nil } - if rollbackErr := h.rollbackQueuedAttachmentClaim(admittedCtx, req.SessionID, source.ID); rollbackErr != nil { + if _, rollbackErr := h.rollbackQueuedAttachmentClaim(admittedCtx, req.SessionID, source.ID); rollbackErr != nil { h.logger.Error("failed to roll back queued attachment", zap.Error(rollbackErr)) return fmt.Errorf("%w: %v", errQueuedAttachmentRollback, rollbackErr) } @@ -433,11 +727,139 @@ func (h *QueueHandlers) claimQueuedAttachmentsAfterInsert(req *wsQueueMessageReq } } +type wsQueueEditRequest struct { + SessionID string `json:"session_id"` + EntryID string `json:"entry_id"` + LeaseID string `json:"lease_id"` + DispatchIfAutoRun bool `json:"dispatch_if_auto_run,omitempty"` +} + +func (h *QueueHandlers) queueEditError(msg *ws.Message, err error) *ws.Message { + switch { + case errors.Is(err, messagequeue.ErrEntryNotFound): + response, _ := ws.NewError(msg.ID, msg.Action, queueErrorCodeEntryNotFound, + "Queue entry was already drained or is no longer editable", nil) + return response + case errors.Is(err, messagequeue.ErrEditLeaseNotFound): + response, _ := ws.NewError(msg.ID, msg.Action, queueErrorCodeEntryNotFound, + "Queue entry was already drained or is no longer editable", nil) + return response + case errors.Is(err, messagequeue.ErrEditConflict): + response, _ := ws.NewError(msg.ID, msg.Action, "edit_conflict", + "Queue entry is being edited by another view", nil) + return response + case errors.Is(err, messagequeue.ErrEditRevisionConflict): + response, _ := ws.NewError(msg.ID, msg.Action, "queue_conflict", + "Queue entry changed while it was being edited", nil) + return response + default: + h.logger.Error("failed to process queue edit", zap.Error(err)) + response, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, + "Failed to process queue edit", nil) + return response + } +} +func (h *QueueHandlers) queueEditLeaseError(msg *ws.Message, err error) *ws.Message { + if errors.Is(err, messagequeue.ErrEditLeaseNotFound) { + response, _ := ws.NewError(msg.ID, msg.Action, "edit_conflict", "Queue entry edit lease is no longer valid", nil) + return response + } + return h.queueEditError(msg, err) +} + +func (h *QueueHandlers) wsBeginEdit(ctx context.Context, msg *ws.Message) (*ws.Message, error) { + var req wsQueueEditRequest + if err := msg.ParsePayload(&req); err != nil { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeBadRequest, "Invalid payload: "+err.Error(), nil) + } + if req.SessionID == "" || req.EntryID == "" { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "session_id and entry_id are required", nil) + } + if denied := h.authorizeSession(ctx, msg, req.SessionID); denied != nil { + return denied, nil + } + if h.queueEdit == nil || ws.ConnectionID(ctx) == "" { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeForbidden, "Queue editing requires a WebSocket connection", nil) + } + lease, err := h.queueEdit.BeginEdit(ctx, req.SessionID, req.EntryID, ws.ConnectionID(ctx)) + if err != nil { + return h.queueEditError(msg, err), nil + } + return ws.NewResponse(msg.ID, msg.Action, lease) +} + +func (h *QueueHandlers) wsRenewEdit(ctx context.Context, msg *ws.Message) (*ws.Message, error) { + var req wsQueueEditRequest + if err := msg.ParsePayload(&req); err != nil { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeBadRequest, "Invalid payload: "+err.Error(), nil) + } + if req.SessionID == "" || req.EntryID == "" || req.LeaseID == "" { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "session_id, entry_id, and lease_id are required", nil) + } + if denied := h.authorizeSession(ctx, msg, req.SessionID); denied != nil { + return denied, nil + } + if h.queueEdit == nil || ws.ConnectionID(ctx) == "" { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeForbidden, "Queue editing requires a WebSocket connection", nil) + } + lease, err := h.queueEdit.RenewEdit(ctx, req.SessionID, req.EntryID, req.LeaseID, ws.ConnectionID(ctx)) + if err != nil { + return h.queueEditLeaseError(msg, err), nil + } + return ws.NewResponse(msg.ID, msg.Action, lease) +} + +func (h *QueueHandlers) wsEndEdit(ctx context.Context, msg *ws.Message) (*ws.Message, error) { + var req wsQueueEditRequest + if err := msg.ParsePayload(&req); err != nil { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeBadRequest, "Invalid payload: "+err.Error(), nil) + } + if req.SessionID == "" || req.EntryID == "" || req.LeaseID == "" { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "session_id, entry_id, and lease_id are required", nil) + } + if denied := h.authorizeSession(ctx, msg, req.SessionID); denied != nil { + return denied, nil + } + if h.queueEdit == nil || ws.ConnectionID(ctx) == "" { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeForbidden, "Queue editing requires a WebSocket connection", nil) + } + var saved bool + var err error + if req.DispatchIfAutoRun { + controller, ok := h.queueEdit.(queueEditSavedLeaseController) + if !ok { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "Queue edit save completion is unavailable", nil) + } + saved, err = controller.EndEditAfterSave(ctx, req.SessionID, req.EntryID, req.LeaseID, ws.ConnectionID(ctx)) + } else { + err = h.queueEdit.EndEdit(ctx, req.SessionID, req.EntryID, req.LeaseID, ws.ConnectionID(ctx)) + } + if err != nil { + return h.queueEditLeaseError(msg, err), nil + } + h.signalPendingAttachmentCleanup(ctx, req.SessionID, req.EntryID) + if saved { + if drainer, ok := h.queueDrainer.(QueueAutoRunPreservingDrainer); ok { + if _, drainErr := drainer.DrainQueuedMessageIfAutoRun(ctx, req.SessionID); drainErr != nil { + h.logger.Warn("post-save queued message drain failed", + zap.String(fieldSessionID, req.SessionID), zap.String(fieldEntryID, req.EntryID), zap.Error(drainErr)) + } + } else { + h.logger.Warn("post-save queued message drain unavailable", + zap.String(fieldSessionID, req.SessionID), zap.String(fieldEntryID, req.EntryID)) + } + } + return ws.NewResponse(msg.ID, msg.Action, map[string]string{fieldSessionID: req.SessionID, fieldEntryID: req.EntryID}) +} + type wsUpdateMessageRequest struct { SessionID string `json:"session_id"` TaskID string `json:"task_id"` SessionIncarnationID string `json:"session_incarnation_id"` EntryID string `json:"entry_id"` + LeaseID string `json:"lease_id,omitempty"` + OperationID string `json:"operation_id,omitempty"` + ExpectedRevision *int64 `json:"expected_target_revision,omitempty"` Content string `json:"content"` Attachments []messagequeue.MessageAttachment `json:"attachments,omitempty"` EntityReferences []v1.EntityReference `json:"entity_references,omitempty"` @@ -527,15 +949,20 @@ func (h *QueueHandlers) wsUpdateMessage(ctx context.Context, msg *ws.Message) (* if messagequeue.IsReservedQueuedBy(req.UserID) { return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, reservedIdentityError(req.UserID), nil) } + connectionID := ws.ConnectionID(ctx) + if connectionID != "" && (h.queueEdit == nil || req.LeaseID == "" || req.OperationID == "" || req.ExpectedRevision == nil) { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, + "lease_id, operation_id, and expected_target_revision are required", nil) + } referencesProvided := req.EntityReferences != nil references, err := h.validateSubmittedReferences(ctx, req.SessionID, "", req.EntityReferences) if err != nil { return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, queueInvalidReferences, nil) } - req.EntityReferences = references // Default empty user_id to QueuedByUser so the UpdateContent guard always // runs against a non-empty owner. Agent entries (queued_by="agent") then // fail the filter, mirroring the canEdit UI gate at the WS layer. + req.EntityReferences = references queuedBy := req.UserID if queuedBy == "" { queuedBy = messagequeue.QueuedByUser @@ -576,42 +1003,688 @@ func (h *QueueHandlers) wsUpdateMessage(ctx context.Context, msg *ws.Message) (* return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, err.Error(), nil) } newlyAdded = newlyAddedQueueAttachments(previous.Attachments, req.Attachments) - if preparer, ok := h.attachmentClaimer.(QueueAttachmentClaimPreparer); ok && h.requiresQueueIdentity() { - prepared, prepareErr := preparer.PrepareQueueAttachmentClaim(ctx, previous.TaskID, queueAttachmentsToV1(newlyAdded)) - if prepareErr != nil { - return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "Attachment is no longer available", nil) - } - atomicClaim = &prepared - } else if err := h.attachmentClaimer.ClaimMessageAttachments(ctx, previous.TaskID, req.SessionID, queueAttachmentsToV1(newlyAdded)); err != nil { - return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "Attachment is no longer available", nil) - } releaseClaims, _ = h.attachmentClaimer.(QueueAttachmentReleaser) - } - updateErr := h.updateQueuedMessage(ctx, identity, req, metadataUpdates, queuedBy, atomicClaim) - if updateErr != nil { - if atomicClaim == nil { - h.releaseFailedQueueAttachmentClaims(ctx, releaseClaims, previous, req.SessionID, newlyAdded) - } - if errors.Is(updateErr, messagequeue.ErrEntryNotFound) { - return ws.NewError(msg.ID, msg.Action, queueErrorCodeEntryNotFound, "Queue entry was already drained or not owned by caller", nil) - } - if isQueueIdentityError(updateErr) { - return queueAccessDeniedResponse(msg), nil - } - return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, updateErr.Error(), nil) - } - if releaseClaims != nil && previous != nil { - if superseded := supersededQueueAttachments(previous.Attachments, req.Attachments); len(superseded) > 0 { - if err := releaseClaims.ReleaseMessageAttachments(ctx, previous.TaskID, req.SessionID, queueAttachmentsToV1(superseded)); err != nil { - h.logger.Warn("failed to release superseded queue attachments", zap.Error(err)) + if h.requiresQueueIdentity() { + if preparer, ok := h.attachmentClaimer.(QueueAttachmentClaimPreparer); ok { + prepared, prepareErr := preparer.PrepareQueueAttachmentClaim(ctx, previous.TaskID, queueAttachmentsToV1(newlyAdded)) + if prepareErr != nil { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "Attachment is no longer available", nil) + } + atomicClaim = &prepared } } } - h.publishStatusForIdentity(ctx, identity) - return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{fieldEntryID: req.EntryID}) + var supersededPending *pendingQueueAttachmentCleanup + if releaseClaims != nil && previous != nil { + cleanupReq := req + cleanupReq.OperationID = attachmentCleanupOperationID(req.OperationID, "superseded") + var prepareErr error + supersededPending, prepareErr = h.preparePendingAttachmentCleanup( + ctx, + cleanupReq, + previous.TaskID, + supersededQueueAttachments(previous.Attachments, req.Attachments), + releaseClaims, + ) + if prepareErr != nil { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, prepareErr.Error(), nil) + } + } + var revision int64 + var updateErr error + applyUpdate := func(updateCtx context.Context) (int64, error) { + attachmentsToClaim := newlyAdded + newlyAdded = nil + var rollbackPending *pendingQueueAttachmentCleanup + if atomicClaim == nil && h.attachmentClaimer != nil { + rollbackReq := req + rollbackReq.OperationID = attachmentCleanupOperationID(req.OperationID, "claim-rollback") + var prepareErr error + rollbackPending, prepareErr = h.preparePendingAttachmentCleanup( + updateCtx, rollbackReq, previous.TaskID, attachmentsToClaim, releaseClaims, + ) + if prepareErr != nil { + return 0, prepareErr + } + if err := h.attachmentClaimer.ClaimMessageAttachments(updateCtx, previous.TaskID, req.SessionID, queueAttachmentsToV1(attachmentsToClaim)); err != nil { + releaseErr := h.releaseQueuedAttachmentUpdateFailure( + updateCtx, previous, req.SessionID, attachmentsToClaim, releaseClaims, + ) + h.settlePendingAttachmentCleanup(rollbackPending, releaseErr) + return 0, fmt.Errorf("%w: %v", errQueuedAttachmentUnavailable, err) + } + } + switch { + case connectionID != "": + revision, updateErr = h.queueEdit.UpdateMessageWithLease(updateCtx, req.SessionID, req.EntryID, + req.LeaseID, req.OperationID, connectionID, *req.ExpectedRevision, req.Content, + req.Attachments, metadataUpdates) + case h.requiresQueueIdentity(): + updateErr = h.updateQueuedMessage(updateCtx, identity, req, metadataUpdates, queuedBy, atomicClaim) + default: + // Non-WebSocket/MCP updates still use the service's durable lease fencing. + updateErr = h.queueService.UpdateMessageWithMetadata(updateCtx, req.SessionID, req.EntryID, + req.Content, req.Attachments, metadataUpdates, queuedBy) + } + if updateErr != nil && atomicClaim == nil && h.attachmentClaimer != nil { + releaseErr := h.releaseQueuedAttachmentUpdateFailure( + updateCtx, previous, req.SessionID, attachmentsToClaim, releaseClaims, + ) + h.settlePendingAttachmentCleanup(rollbackPending, releaseErr) + } else { + h.settlePendingAttachmentCleanup(rollbackPending, nil) + } + return revision, updateErr + } + if connectionID != "" && h.attachmentClaimer != nil { + revision, updateErr = h.updateMessageWithAttachmentLease( + ctx, req, connectionID, previous, releaseClaims, supersededPending, + &newlyAdded, metadataUpdates, applyUpdate, + ) + } else { + revision, updateErr = applyUpdate(ctx) + } + if updateErr != nil { + h.settlePendingAttachmentCleanup(supersededPending, nil) + return h.queueUpdateFailure(ctx, msg, req, updateErr) + } + if releaseClaims != nil && supersededPending != nil && connectionID == "" { + releaseErr := h.releaseSupersededQueueAttachments( + ctx, supersededPending.req, supersededPending.previous, releaseClaims, + ) + h.settlePendingAttachmentCleanup(supersededPending, releaseErr) + } + response := map[string]interface{}{fieldEntryID: req.EntryID} + if req.OperationID != "" { + response["operation_id"] = req.OperationID + response["target_revision"] = revision + } + h.publishStatus(ctx, req.SessionID) + return ws.NewResponse(msg.ID, msg.Action, response) +} + +func (h *QueueHandlers) releaseQueuedAttachmentUpdateFailure( + ctx context.Context, + previous *messagequeue.QueuedMessage, + sessionID string, + newlyAdded []messagequeue.MessageAttachment, + releaseClaims QueueAttachmentReleaser, +) error { + if releaseClaims == nil || previous == nil || len(newlyAdded) == 0 { + return nil + } + cleanupCtx := context.WithoutCancel(ctx) + candidates, err := h.unreferencedQueueAttachments(cleanupCtx, sessionID, previous.ID, newlyAdded) + if err != nil { + return err + } + if len(candidates) == 0 { + return nil + } + if releaseErr := releaseClaims.ReleaseMessageAttachments(cleanupCtx, previous.TaskID, sessionID, queueAttachmentsToV1(candidates)); releaseErr != nil { + h.logger.Warn("failed to release attachments after queue update failure", zap.Error(releaseErr)) + return releaseErr + } + return nil +} + +// unreferencedQueueAttachments prevents cleanup from deleting a claim that a +// different pending entry in the same session still uses. A reference lookup +// failure fails closed because releasing in that case can destroy another +// queued prompt's attachment. +func (h *QueueHandlers) unreferencedQueueAttachments( + ctx context.Context, + sessionID, excludedEntryID string, + candidates []messagequeue.MessageAttachment, +) ([]messagequeue.MessageAttachment, error) { + ctx = context.WithoutCancel(ctx) + if len(candidates) == 0 { + return nil, nil + } + unique := make([]messagequeue.MessageAttachment, 0, len(candidates)) + seen := make(map[string]struct{}, len(candidates)) + for _, attachment := range candidates { + if attachment.AttachmentID == "" { + continue + } + if _, ok := seen[attachment.AttachmentID]; ok { + continue + } + seen[attachment.AttachmentID] = struct{}{} + unique = append(unique, attachment) + } + if len(unique) == 0 { + return nil, nil + } + checker, ok := h.queueService.(queueAttachmentReferenceChecker) + if !ok { + return unique, nil + } + ids := make([]string, 0, len(unique)) + for _, attachment := range unique { + ids = append(ids, attachment.AttachmentID) + } + referenced, err := checker.ReferencedQueueAttachmentIDs(ctx, sessionID, excludedEntryID, ids) + if err != nil { + h.logger.Warn("failed to inspect queue attachment references", zap.Error(err)) + return nil, err + } + unreferenced := make([]messagequeue.MessageAttachment, 0, len(unique)) + for _, attachment := range unique { + if _, ok := referenced[attachment.AttachmentID]; !ok { + unreferenced = append(unreferenced, attachment) + } + } + return unreferenced, nil +} + +// releaseSupersededQueueAttachments serializes attachment cleanup with later +// edits and recomputes the retained set from the current queue row. A second +// edit may reclaim an attachment after this update commits but before its +// cleanup runs; releasing from the stale request snapshot would then destroy +// a claim still referenced by the row. +func (h *QueueHandlers) releaseSupersededQueueAttachments( + ctx context.Context, + req wsUpdateMessageRequest, + previous *messagequeue.QueuedMessage, + releaser QueueAttachmentReleaser, +) error { + release := func(admittedCtx context.Context) error { + return h.releaseSupersededQueueAttachmentsAdmitted(admittedCtx, req, previous, releaser) + } + if admission, ok := h.queueService.(queueEditAdmissionController); ok { + if err := admission.WithSessionAdmission(context.WithoutCancel(ctx), req.SessionID, release); err != nil { + h.logger.Warn("failed to serialize superseded queue attachment cleanup", zap.Error(err)) + return err + } + return nil + } + if err := release(context.WithoutCancel(ctx)); err != nil { + h.logger.Warn("failed to release superseded queue attachments", zap.Error(err)) + return err + } + return nil +} + +func (h *QueueHandlers) releaseSupersededQueueAttachmentsAdmitted( + ctx context.Context, + req wsUpdateMessageRequest, + previous *messagequeue.QueuedMessage, + releaser QueueAttachmentReleaser, +) error { + cleanupCtx := context.WithoutCancel(ctx) + current, err := h.queueService.GetEntry(cleanupCtx, req.SessionID, req.EntryID) + if errors.Is(err, messagequeue.ErrEntryNotFound) { + if locator, ok := h.queueService.(queueEntryLocator); ok { + current, err = locator.FindEntryByID(cleanupCtx, req.EntryID) + if err == nil && current != nil { + req.SessionID = current.SessionID + } + } + if errors.Is(err, messagequeue.ErrEntryNotFound) { + return h.releaseQueueAttachmentCandidates( + cleanupCtx, previous.TaskID, req, previous.Attachments, releaser, + ) + } + } + if err != nil { + h.logger.Warn("failed to reload queue entry before attachment cleanup", zap.Error(err)) + return err + } + if current == nil { + return nil + } + return h.releaseQueueAttachmentCandidates( + cleanupCtx, previous.TaskID, req, supersededQueueAttachments(previous.Attachments, current.Attachments), releaser, + ) +} + +func (h *QueueHandlers) releaseQueueAttachmentCandidates( + ctx context.Context, + taskID string, + req wsUpdateMessageRequest, + candidates []messagequeue.MessageAttachment, + releaser QueueAttachmentReleaser, +) error { + cleanupCtx := context.WithoutCancel(ctx) + unreferenced, err := h.unreferencedQueueAttachments(cleanupCtx, req.SessionID, req.EntryID, candidates) + if err != nil { + return err + } + if len(unreferenced) == 0 { + return nil + } + if err := releaser.ReleaseMessageAttachments( + cleanupCtx, taskID, req.SessionID, queueAttachmentsToV1(unreferenced), + ); err != nil { + h.logger.Warn("failed to release superseded queue attachments", zap.Error(err)) + return err + } + return nil +} + +func (h *QueueHandlers) preparePendingAttachmentCleanup( + ctx context.Context, + req wsUpdateMessageRequest, + taskID string, + attachments []messagequeue.MessageAttachment, + releaser QueueAttachmentReleaser, +) (*pendingQueueAttachmentCleanup, error) { + return h.preparePendingAttachmentCleanupWithState( + ctx, req, taskID, attachments, releaser, false, nil, + ) +} + +func (h *QueueHandlers) preparePendingAttachmentCleanupWithState( + ctx context.Context, + req wsUpdateMessageRequest, + taskID string, + attachments []messagequeue.MessageAttachment, + releaser QueueAttachmentReleaser, + claimPending bool, + expectedEntry *messagequeue.QueuedMessage, +) (*pendingQueueAttachmentCleanup, error) { + if releaser == nil || len(attachments) == 0 { + return nil, nil + } + entryFingerprint, err := queuedMessageFingerprint(expectedEntry) + if err != nil { + return nil, err + } + key := pendingQueueAttachmentCleanupKey{ + sessionID: req.SessionID, entryID: req.EntryID, operationID: req.OperationID, + } + cleanupCtx := context.Background() + if identity, ok := authn.IdentityFromContext(ctx); ok { + cleanupCtx = authn.WithIdentity(cleanupCtx, identity) + } + previous := &messagequeue.QueuedMessage{ + ID: req.EntryID, SessionID: req.SessionID, TaskID: taskID, + Attachments: append([]messagequeue.MessageAttachment(nil), attachments...), + } + pending := &pendingQueueAttachmentCleanup{ + key: key, + req: req, + previous: previous, + releaser: releaser, removeEntry: claimPending, claimPending: claimPending, + entryFingerprint: entryFingerprint, + authCtx: cleanupCtx, wake: make(chan struct{}, 1), + } + persist := func(admittedCtx context.Context) error { + if _, ok := h.attachmentCleanupStore(); !ok { + return nil + } + current, err := h.queueService.GetEntry(admittedCtx, req.SessionID, req.EntryID) + if err != nil { + return fmt.Errorf("reload queue entry before cleanup persistence: %w", err) + } + h.setPendingAttachmentCleanupSessionID(pending, current.SessionID) + return h.persistPendingAttachmentCleanup(admittedCtx, pending) + } + if admission, ok := h.queueService.(queueEditAdmissionController); ok { + admissionCtx := ctx + if admissionCtx.Err() != nil { + admissionCtx = cleanupCtx + } + if err := admission.WithSessionAdmission(admissionCtx, req.SessionID, persist); err != nil { + return nil, err + } + return pending, nil + } + if err := persist(cleanupCtx); err != nil { + return nil, err + } + return pending, nil +} + +func (h *QueueHandlers) persistPendingAttachmentCleanup( + ctx context.Context, + pending *pendingQueueAttachmentCleanup, +) error { + store, ok := h.attachmentCleanupStore() + if !ok { + return nil + } + identity, hasIdentity := authn.IdentityFromContext(pending.authCtx) + if !hasIdentity || identity.UserID == "" { + return errors.New("queue attachment cleanup requires an owner identity") + } + if err := store.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: pending.req.SessionID, CurrentSessionID: h.pendingAttachmentCleanupSessionID(pending), + EntryID: pending.req.EntryID, OperationID: pending.req.OperationID, + TaskID: pending.previous.TaskID, OwnerID: identity.UserID, LeaseID: pending.req.LeaseID, + RemoveEntry: pending.removeEntry, + ClaimPending: pending.claimPending, + EntryFingerprint: pending.entryFingerprint, + Attachments: append([]messagequeue.MessageAttachment(nil), pending.previous.Attachments...), + CreatedAt: time.Now().UTC(), + }); err != nil { + return fmt.Errorf("persist queue attachment cleanup: %w", err) + } + return nil +} + +func (h *QueueHandlers) queuePendingAttachmentCleanup(pending *pendingQueueAttachmentCleanup) { + if pending == nil { + return + } + h.attachmentCleanupMu.Lock() + if h.attachmentCleanupStopped { + h.attachmentCleanupMu.Unlock() + return + } + if _, exists := h.pendingAttachmentCleanup[pending.key]; exists { + h.attachmentCleanupMu.Unlock() + return + } + h.pendingAttachmentCleanup[pending.key] = pending + if !h.attachmentCleanupStarted { + h.attachmentCleanupMu.Unlock() + return + } + h.attachmentCleanupWG.Add(1) + h.attachmentCleanupMu.Unlock() + go h.retryPendingAttachmentCleanup(pending) +} + +func (h *QueueHandlers) settlePendingAttachmentCleanup( + pending *pendingQueueAttachmentCleanup, + cleanupErr error, +) { + if pending == nil { + return + } + if cleanupErr != nil { + h.queuePendingAttachmentCleanup(pending) + return + } + if err := h.deletePendingAttachmentCleanup(pending); err != nil { + h.logger.Warn("failed to acknowledge queue attachment cleanup", zap.Error(err)) + h.queuePendingAttachmentCleanup(pending) + } +} + +func (h *QueueHandlers) retryPendingAttachmentCleanup(pending *pendingQueueAttachmentCleanup) { + defer h.attachmentCleanupWG.Done() + delay := 10 * time.Millisecond + for { + if !h.editLeaseActive(pending) { + if err := h.runPendingAttachmentCleanup(pending); err == nil { + if err := h.deletePendingAttachmentCleanup(pending); err != nil { + h.logger.Warn("failed to acknowledge queue attachment cleanup", zap.Error(err)) + } else { + h.attachmentCleanupMu.Lock() + delete(h.pendingAttachmentCleanup, pending.key) + h.attachmentCleanupMu.Unlock() + return + } + } + } + if !h.waitForAttachmentCleanupRetry(delay, pending.wake) { + return + } + if delay < time.Second { + delay *= 2 + } + } +} + +func (h *QueueHandlers) deletePendingAttachmentCleanup(pending *pendingQueueAttachmentCleanup) error { + store, ok := h.attachmentCleanupStore() + if !ok { + return nil + } + return store.DeleteAttachmentCleanup( + pending.authCtx, pending.key.sessionID, pending.key.entryID, pending.key.operationID, + ) +} + +func (h *QueueHandlers) editLeaseActive(pending *pendingQueueAttachmentCleanup) bool { + reader, ok := h.queueService.(queueEditLeaseStateReader) + if !ok { + return false + } + lease, err := reader.GetEditLease( + pending.authCtx, h.pendingAttachmentCleanupSessionID(pending), pending.req.EntryID, + ) + if err != nil { + return !errors.Is(err, messagequeue.ErrEditLeaseNotFound) + } + return lease != nil +} + +func (h *QueueHandlers) waitForAttachmentCleanupRetry(delay time.Duration, wake <-chan struct{}) bool { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-h.attachmentCleanupCtx.Done(): + return false + case <-wake: + return true + case <-timer.C: + return true + } +} + +func (h *QueueHandlers) signalPendingAttachmentCleanup(ctx context.Context, sessionID, entryID string) { + h.attachmentCleanupMu.Lock() + pendingCleanups := make([]*pendingQueueAttachmentCleanup, 0) + for key, pending := range h.pendingAttachmentCleanup { + if key.entryID == entryID { + pendingCleanups = append(pendingCleanups, pending) + } + } + h.attachmentCleanupMu.Unlock() + for _, pending := range pendingCleanups { + if err := h.refreshPendingAttachmentCleanupSessionID(ctx, pending); err != nil { + h.logger.Warn("failed to refresh queue attachment cleanup before wake", zap.Error(err)) + } + currentSessionID := h.pendingAttachmentCleanupSessionID(pending) + if pending.key.sessionID != sessionID && currentSessionID != sessionID { + continue + } + select { + case pending.wake <- struct{}{}: + default: + } + } +} + +func (h *QueueHandlers) pendingAttachmentCleanupSessionID( + pending *pendingQueueAttachmentCleanup, +) string { + h.attachmentCleanupMu.Lock() + defer h.attachmentCleanupMu.Unlock() + if pending.currentSessionID != "" { + return pending.currentSessionID + } + return pending.req.SessionID +} + +func (h *QueueHandlers) setPendingAttachmentCleanupSessionID( + pending *pendingQueueAttachmentCleanup, + sessionID string, +) { + h.attachmentCleanupMu.Lock() + pending.currentSessionID = sessionID + h.attachmentCleanupMu.Unlock() +} + +func (h *QueueHandlers) refreshPendingAttachmentCleanupSessionID( + ctx context.Context, + pending *pendingQueueAttachmentCleanup, +) error { + if _, ok := h.attachmentCleanupStore(); !ok { + return nil + } + locator, ok := h.queueService.(queueAttachmentCleanupLocator) + if !ok { + return nil + } + cleanup, err := locator.GetAttachmentCleanup( + ctx, pending.key.sessionID, pending.key.entryID, pending.key.operationID, + ) + if err != nil { + return err + } + if cleanup == nil { + return nil + } + sessionID := cleanup.CurrentSessionID + if sessionID == "" { + sessionID = cleanup.SessionID + } + h.setPendingAttachmentCleanupSessionID(pending, sessionID) + return nil +} + +func (h *QueueHandlers) runPendingAttachmentCleanup(pending *pendingQueueAttachmentCleanup) error { + if err := h.refreshPendingAttachmentCleanupSessionID(pending.authCtx, pending); err != nil { + return err + } + if _, err := h.retargetPendingAttachmentCleanup(pending.authCtx, pending); err != nil { + return err + } + sessionID := h.pendingAttachmentCleanupSessionID(pending) + cleanup := func(ctx context.Context) error { + if h.editLeaseActive(pending) { + return errAttachmentCleanupLeaseActive + } + return h.settlePendingAttachmentCleanupAdmitted(ctx, pending) + } + if admission, ok := h.queueService.(queueEditAdmissionController); ok { + return admission.WithSessionAdmission(pending.authCtx, sessionID, cleanup) + } + return cleanup(pending.authCtx) +} + +func (h *QueueHandlers) retargetPendingAttachmentCleanup( + ctx context.Context, + pending *pendingQueueAttachmentCleanup, +) (bool, error) { + sessionID := h.pendingAttachmentCleanupSessionID(pending) + current, err := h.queueService.GetEntry(ctx, sessionID, pending.req.EntryID) + if err == nil { + return current != nil, nil + } + if !errors.Is(err, messagequeue.ErrEntryNotFound) { + return false, err + } + locator, ok := h.queueService.(queueEntryLocator) + if !ok { + return false, nil + } + current, err = locator.FindEntryByID(ctx, pending.req.EntryID) + if errors.Is(err, messagequeue.ErrEntryNotFound) { + return false, nil + } + if err != nil { + return false, err + } + if current == nil { + return false, nil + } + h.setPendingAttachmentCleanupSessionID(pending, current.SessionID) + return true, nil +} + +func (h *QueueHandlers) settlePendingAttachmentCleanupAdmitted( + ctx context.Context, + pending *pendingQueueAttachmentCleanup, +) error { + sessionID := h.pendingAttachmentCleanupSessionID(pending) + req := pending.req + req.SessionID = sessionID + current, err := h.queueService.GetEntry(ctx, sessionID, req.EntryID) + if errors.Is(err, messagequeue.ErrEntryNotFound) { + found, retargetErr := h.retargetPendingAttachmentCleanup(ctx, pending) + if retargetErr != nil { + return retargetErr + } + if found { + return errAttachmentCleanupEntryChanged + } + return h.releaseQueueAttachmentCandidates( + ctx, pending.previous.TaskID, req, pending.previous.Attachments, pending.releaser, + ) + } + if err != nil { + return err + } + if !pending.removeEntry && !pending.claimPending { + return h.releaseQueueAttachmentCandidates( + ctx, + pending.previous.TaskID, + req, + supersededQueueAttachments(pending.previous.Attachments, current.Attachments), + pending.releaser, + ) + } + currentFingerprint, err := queuedMessageFingerprint(current) + if err != nil { + return err + } + if pending.entryFingerprint == "" || currentFingerprint != pending.entryFingerprint { + if sessionID != pending.key.sessionID { + return errAttachmentCleanupEntryChanged + } + return h.releaseQueueAttachmentCandidates( + ctx, + pending.previous.TaskID, + req, + supersededQueueAttachments(pending.previous.Attachments, current.Attachments), + pending.releaser, + ) + } + if pending.removeEntry { + if _, err := h.rollbackQueuedAttachmentClaim(ctx, sessionID, req.EntryID); err != nil { + return err + } + return h.releaseQueueAttachmentCandidates( + ctx, pending.previous.TaskID, req, pending.previous.Attachments, pending.releaser, + ) + } + return h.attachmentClaimer.ClaimMessageAttachments( + context.WithoutCancel(ctx), + pending.previous.TaskID, + sessionID, + queueAttachmentsToV1(pending.previous.Attachments), + ) +} + +func queuedMessageFingerprint(msg *messagequeue.QueuedMessage) (string, error) { + if msg == nil { + return "", nil + } + attachments := msg.Attachments + if attachments == nil { + attachments = []messagequeue.MessageAttachment{} + } + metadata := msg.Metadata + if metadata == nil { + metadata = map[string]interface{}{} + } + encoded, err := json.Marshal(struct { + ID string `json:"id"` + TaskID string `json:"task_id"` + Content string `json:"content"` + Model string `json:"model"` + PlanMode bool `json:"plan_mode"` + Attachments []messagequeue.MessageAttachment `json:"attachments"` + Metadata map[string]interface{} `json:"metadata"` + QueuedAt time.Time `json:"queued_at"` + QueuedBy string `json:"queued_by"` + }{ + ID: msg.ID, TaskID: msg.TaskID, Content: msg.Content, Model: msg.Model, PlanMode: msg.PlanMode, + Attachments: attachments, Metadata: metadata, QueuedAt: msg.QueuedAt.UTC(), QueuedBy: msg.QueuedBy, + }) + if err != nil { + return "", fmt.Errorf("fingerprint queued attachment source: %w", err) + } + sum := sha256.Sum256(encoded) + return hex.EncodeToString(sum[:]), nil } -// newlyAddedQueueAttachments returns attachment descriptors that were not present before. func newlyAddedQueueAttachments(previous, replacement []messagequeue.MessageAttachment) []messagequeue.MessageAttachment { retained := make(map[string]struct{}, len(previous)) for _, attachment := range previous { @@ -624,12 +1697,145 @@ func newlyAddedQueueAttachments(previous, replacement []messagequeue.MessageAtta if attachment.AttachmentID == "" { continue } + if _, ok := retained[attachment.AttachmentID]; !ok { newlyAdded = append(newlyAdded, attachment) } } return newlyAdded } +func attachmentCleanupOperationID(operationID, kind string) string { + if operationID == "" { + return kind + "-" + uuid.NewString() + } + return operationID + ":" + kind +} + +func (h *QueueHandlers) updateMessageWithAttachmentLease( + ctx context.Context, + req wsUpdateMessageRequest, + connectionID string, + previous *messagequeue.QueuedMessage, + releaseClaims QueueAttachmentReleaser, + supersededPending *pendingQueueAttachmentCleanup, + newlyAdded *[]messagequeue.MessageAttachment, + metadataUpdates map[string]interface{}, + applyUpdate func(context.Context) (int64, error), +) (int64, error) { + controller, ok := h.queueEdit.(queueEditAttachmentController) + if !ok { + return h.updateMessageWithAttachmentAdmissionFallback( + ctx, req, releaseClaims, supersededPending, applyUpdate, + ) + } + + var claimedAttachments []messagequeue.MessageAttachment + var claimAttempted bool + var rollbackPending *pendingQueueAttachmentCleanup + return controller.UpdateMessageWithLeaseAfterValidationAndFinalize( + ctx, req.SessionID, req.EntryID, req.LeaseID, req.OperationID, connectionID, + *req.ExpectedRevision, req.Content, req.Attachments, metadataUpdates, + func(prepareCtx context.Context) error { + claimedAttachments = append(claimedAttachments[:0], *newlyAdded...) + *newlyAdded = nil + rollbackReq := req + rollbackReq.OperationID = attachmentCleanupOperationID(req.OperationID, "claim-rollback") + var err error + rollbackPending, err = h.preparePendingAttachmentCleanup( + prepareCtx, rollbackReq, previous.TaskID, claimedAttachments, releaseClaims, + ) + if err != nil { + return err + } + claimAttempted = true + if err := h.attachmentClaimer.ClaimMessageAttachments(prepareCtx, previous.TaskID, req.SessionID, queueAttachmentsToV1(claimedAttachments)); err != nil { + return fmt.Errorf("%w: %v", errQueuedAttachmentUnavailable, err) + } + return nil + }, + func(rollbackCtx context.Context) error { + if claimAttempted && releaseClaims != nil && previous != nil { + if len(claimedAttachments) == 0 { + // Invoke the release callback after every attempted + // claim, including a no-op claim, so rollback ordering + // remains observable to attachment lifecycle owners. + releaseErr := releaseClaims.ReleaseMessageAttachments( + context.WithoutCancel(rollbackCtx), previous.TaskID, req.SessionID, nil, + ) + h.settlePendingAttachmentCleanup(rollbackPending, releaseErr) + } else { + releaseErr := h.releaseQueuedAttachmentUpdateFailure( + rollbackCtx, previous, req.SessionID, claimedAttachments, releaseClaims, + ) + h.settlePendingAttachmentCleanup(rollbackPending, releaseErr) + } + } + claimedAttachments = nil + rollbackPending = nil + return nil + }, + func(finalizeCtx context.Context, _ *messagequeue.QueuedMessage) error { + h.settlePendingAttachmentCleanup(rollbackPending, nil) + rollbackPending = nil + if releaseClaims == nil || supersededPending == nil { + return nil + } + releaseErr := h.releaseSupersededQueueAttachmentsAdmitted( + finalizeCtx, supersededPending.req, supersededPending.previous, releaseClaims, + ) + h.settlePendingAttachmentCleanup(supersededPending, releaseErr) + return nil + }, + ) +} + +func (h *QueueHandlers) updateMessageWithAttachmentAdmissionFallback( + ctx context.Context, + req wsUpdateMessageRequest, + releaseClaims QueueAttachmentReleaser, + supersededPending *pendingQueueAttachmentCleanup, + applyUpdate func(context.Context) (int64, error), +) (int64, error) { + admission, ok := h.queueEdit.(queueEditAdmissionController) + if !ok { + return applyUpdate(ctx) + } + var revision int64 + err := admission.WithSessionAdmission(ctx, req.SessionID, func(admittedCtx context.Context) error { + var err error + revision, err = applyUpdate(admittedCtx) + if err != nil { + return err + } + if releaseClaims != nil && supersededPending != nil { + cleanupErr := h.releaseSupersededQueueAttachmentsAdmitted( + admittedCtx, supersededPending.req, supersededPending.previous, releaseClaims, + ) + h.settlePendingAttachmentCleanup(supersededPending, cleanupErr) + } + return nil + }) + return revision, err +} +func (h *QueueHandlers) queueUpdateFailure( + _ context.Context, + msg *ws.Message, + req wsUpdateMessageRequest, + updateErr error, +) (*ws.Message, error) { + if errors.Is(updateErr, errQueuedAttachmentUnavailable) { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "Attachment is no longer available", nil) + } + if errors.Is(updateErr, messagequeue.ErrEditConflict) || + errors.Is(updateErr, messagequeue.ErrEditLeaseNotFound) || + errors.Is(updateErr, messagequeue.ErrEditRevisionConflict) { + return h.queueEditLeaseError(msg, updateErr), nil + } + if errors.Is(updateErr, messagequeue.ErrEntryNotFound) { + return ws.NewError(msg.ID, msg.Action, queueErrorCodeEntryNotFound, "Queue entry was already drained or not owned by caller", nil) + } + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, updateErr.Error(), nil) +} // supersededQueueAttachments returns attachment descriptors dropped by the replacement. func supersededQueueAttachments(previous, replacement []messagequeue.MessageAttachment) []messagequeue.MessageAttachment { @@ -666,20 +1872,69 @@ func (h *QueueHandlers) validateSubmittedReferences( return h.referenceValidator.ValidateForSubmission(ctx, sessionID, taskID, references) } -// rollbackQueuedAttachmentClaim releases attachments claimed for a queue entry that failed to persist. -func (h *QueueHandlers) rollbackQueuedAttachmentClaim(ctx context.Context, sessionID, entryID string) error { - if err := h.queueService.RemoveEntry(ctx, sessionID, entryID); err == nil { - return nil +// rollbackQueuedAttachmentClaim removes the queue source whose attachment +// claim failed. removed is false only when another mutation already removed it. +func (h *QueueHandlers) rollbackQueuedAttachmentClaim( + ctx context.Context, + sessionID, entryID string, +) (bool, error) { + rollbackCtx := context.WithoutCancel(ctx) + if err := h.queueService.RemoveEntry(rollbackCtx, sessionID, entryID); err == nil { + return true, nil + } else if errors.Is(err, messagequeue.ErrEntryNotFound) { + return false, nil } else { h.logger.Error("failed to remove queue entry after attachment claim failure", zap.String("entry_id", entryID), zap.Error(err)) } taker, ok := h.queueService.(queueEntryTaker) if !ok { - return errors.New("queue service cannot atomically remove a queued entry") + return false, errors.New("queue service cannot atomically remove a queued entry") } - _, _, err := taker.TakeQueuedEntry(ctx, sessionID, entryID) - return err + _, removed, err := taker.TakeQueuedEntry(rollbackCtx, sessionID, entryID) + return removed, err +} + +func (h *QueueHandlers) prepareEntryRemovalCleanup( + ctx context.Context, + entry *messagequeue.QueuedMessage, + operationID string, +) (*pendingQueueAttachmentCleanup, error) { + if entry == nil || len(entry.Attachments) == 0 || h.attachmentClaimer == nil { + return nil, nil + } + releaser, ok := h.attachmentClaimer.(QueueAttachmentReleaser) + if !ok { + return nil, nil + } + return h.preparePendingAttachmentCleanup(ctx, wsUpdateMessageRequest{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: operationID, + }, entry.TaskID, entry.Attachments, releaser) +} + +// releaseQueuedAttachments releases descriptors no longer referenced by the +// remaining queue entries. The queue entry has already been removed. +func (h *QueueHandlers) releaseQueuedAttachments( + ctx context.Context, + entry *messagequeue.QueuedMessage, + pending *pendingQueueAttachmentCleanup, +) { + if entry == nil || pending == nil { + return + } + cleanupCtx := context.WithoutCancel(ctx) + candidates, releaseErr := h.unreferencedQueueAttachments( + cleanupCtx, entry.SessionID, entry.ID, entry.Attachments, + ) + if len(candidates) > 0 { + releaseErr = pending.releaser.ReleaseMessageAttachments( + cleanupCtx, entry.TaskID, entry.SessionID, queueAttachmentsToV1(candidates), + ) + if releaseErr != nil { + h.logger.Warn("failed to release attachments after queue entry removal", zap.Error(releaseErr)) + } + } + h.settlePendingAttachmentCleanup(pending, releaseErr) } // firstInvalidDeliveryMode returns the index of the first attachment with an unknown delivery mode. diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_admission_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_admission_test.go new file mode 100644 index 0000000000..1d4bb88412 --- /dev/null +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_admission_test.go @@ -0,0 +1,784 @@ +package handlers + +import ( + "context" + "errors" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/kandev/kandev/internal/auth/authn" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + v1 "github.com/kandev/kandev/pkg/api/v1" + ws "github.com/kandev/kandev/pkg/websocket" + "github.com/stretchr/testify/require" +) + +type admissionCleanupAssertingClaimer struct { + store queueAttachmentCleanupStore + claimAttempts atomic.Int32 + sawObligation atomic.Bool + releaseAttempt atomic.Int32 +} + +func (c *admissionCleanupAssertingClaimer) ClaimMessageAttachments( + ctx context.Context, + _, _ string, + _ []v1.MessageAttachment, +) error { + c.claimAttempts.Add(1) + cleanups, err := c.store.ListAttachmentCleanups(ctx) + if err == nil && len(cleanups) == 1 && cleanups[0].ClaimPending { + c.sawObligation.Store(true) + } + return errors.New("claim outcome unavailable") +} + +func (c *admissionCleanupAssertingClaimer) ReleaseMessageAttachments( + context.Context, + string, + string, + []v1.MessageAttachment, +) error { + c.releaseAttempt.Add(1) + return nil +} + +func TestQueueAttachmentClaimFailureHasDurableCleanupBeforeClaim(t *testing.T) { + handlers, queue, db := newPersistentCleanupQueue(t, filepath.Join(t.TempDir(), "queue.db")) + defer func() { _ = db.Close() }() + claimer := &admissionCleanupAssertingClaimer{store: queue} + handlers.SetAttachmentClaimer(claimer) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + + response, err := handlers.wsQueueMessage(ctx, createTestMessage(t, ws.ActionMessageQueueAdd, map[string]interface{}{ + "session_id": "session-admission-cleanup", + "task_id": "task-admission-cleanup", + "content": "queued", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-admission-cleanup", + Name: "attachment.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + })) + + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + require.Equal(t, int32(1), claimer.claimAttempts.Load()) + require.True(t, claimer.sawObligation.Load(), "cleanup obligation must exist before the claim starts") + require.Equal(t, int32(1), claimer.releaseAttempt.Load()) + require.Empty(t, queue.GetStatus(ctx, "session-admission-cleanup").Entries) + cleanups, err := queue.ListAttachmentCleanups(ctx) + require.NoError(t, err) + require.Empty(t, cleanups) +} + +func TestQueuedMessageFingerprintSurvivesSQLiteRoundTrip(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + _, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := context.Background() + entry, err := queue.QueueMessage( + ctx, "session-fingerprint", "task-fingerprint", "queued", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{Type: "resource", AttachmentID: "attachment-fingerprint"}}, + ) + require.NoError(t, err) + require.NoError(t, db.Close()) + _, reopenedQueue, reopenedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = reopenedDB.Close() }() + stored, err := reopenedQueue.GetEntry(ctx, entry.SessionID, entry.ID) + require.NoError(t, err) + before, err := queuedMessageFingerprint(entry) + require.NoError(t, err) + after, err := queuedMessageFingerprint(stored) + require.NoError(t, err) + require.Equal(t, before, after, "queued message changed during SQLite restart:\nbefore: %#v\nafter: %#v", entry, stored) +} + +func TestAdmissionAttachmentCleanupRemovesEntryAfterRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + _, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := context.Background() + entry, err := queue.QueueMessage( + ctx, "session-admission-restart", "task-admission-restart", "queued", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-admission-restart", + Name: "attachment.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + fingerprint, err := queuedMessageFingerprint(entry) + require.NoError(t, err) + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "admission-restart", + TaskID: entry.TaskID, OwnerID: "owner", RemoveEntry: true, + EntryFingerprint: fingerprint, Attachments: entry.Attachments, + })) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + reloadedCleanups, err := restartedQueue.ListAttachmentCleanups(ctx) + require.NoError(t, err) + require.Len(t, reloadedCleanups, 1) + require.Equal(t, fingerprint, reloadedCleanups[0].EntryFingerprint) + reloadedEntry, err := restartedQueue.GetEntry(ctx, entry.SessionID, entry.ID) + require.NoError(t, err) + reloadedFingerprint, err := queuedMessageFingerprint(reloadedEntry) + require.NoError(t, err) + require.Equal(t, fingerprint, reloadedFingerprint) + claimer := &controlledCleanupClaimer{} + restarted.SetAttachmentClaimer(claimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + _, getErr := restartedQueue.GetEntry(ctx, entry.SessionID, entry.ID) + cleanups, listErr := restartedQueue.ListAttachmentCleanups(ctx) + return errors.Is(getErr, messagequeue.ErrEntryNotFound) && + claimer.released.Load() == 1 && listErr == nil && len(cleanups) == 0 + }, time.Second, 10*time.Millisecond) + cleanups, err := restartedQueue.ListAttachmentCleanups(ctx) + require.NoError(t, err) + require.Empty(t, cleanups) +} + +type admissionClaimRecoveryClaimer struct { + claims atomic.Int32 + releases atomic.Int32 +} + +func (c *admissionClaimRecoveryClaimer) ClaimMessageAttachments( + context.Context, + string, + string, + []v1.MessageAttachment, +) error { + c.claims.Add(1) + return nil +} + +func (c *admissionClaimRecoveryClaimer) ReleaseMessageAttachments( + context.Context, + string, + string, + []v1.MessageAttachment, +) error { + c.releases.Add(1) + return nil +} + +func TestPendingAdmissionClaimResumesWithoutRemovingEntry(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + _, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := context.Background() + entry, err := queue.QueueMessage( + ctx, "session-claim-restart", "task-claim-restart", "queued", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-claim-restart", + Name: "attachment.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + fingerprint, err := queuedMessageFingerprint(entry) + require.NoError(t, err) + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "claim-restart", + TaskID: entry.TaskID, OwnerID: "owner", ClaimPending: true, + EntryFingerprint: fingerprint, Attachments: entry.Attachments, + })) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + claimer := &admissionClaimRecoveryClaimer{} + restarted.SetAttachmentClaimer(claimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + cleanups, listErr := restartedQueue.ListAttachmentCleanups(ctx) + return listErr == nil && len(cleanups) == 0 && claimer.claims.Load() == 1 + }, time.Second, 10*time.Millisecond) + require.Zero(t, claimer.releases.Load()) + current, err := restartedQueue.GetEntry(ctx, entry.SessionID, entry.ID) + require.NoError(t, err) + require.Equal(t, entry.ID, current.ID) +} + +type failSecondCleanupUpsertQueue struct { + QueueService + service *messagequeue.Service + upserts atomic.Int32 +} + +func (q *failSecondCleanupUpsertQueue) AttachmentCleanupPersistenceAvailable() bool { return true } + +func (q *failSecondCleanupUpsertQueue) UpsertAttachmentCleanup( + ctx context.Context, + cleanup messagequeue.AttachmentCleanup, +) error { + if q.upserts.Add(1) == 2 { + return errors.New("cleanup state update unavailable") + } + return q.service.UpsertAttachmentCleanup(ctx, cleanup) +} + +func (q *failSecondCleanupUpsertQueue) DeleteAttachmentCleanup( + ctx context.Context, + sessionID, entryID, operationID string, +) error { + return q.service.DeleteAttachmentCleanup(ctx, sessionID, entryID, operationID) +} + +func (q *failSecondCleanupUpsertQueue) ListAttachmentCleanups( + ctx context.Context, +) ([]messagequeue.AttachmentCleanup, error) { + return q.service.ListAttachmentCleanups(ctx) +} + +func (q *failSecondCleanupUpsertQueue) WithSessionAdmission( + ctx context.Context, + sessionID string, + fn func(context.Context) error, +) error { + return q.service.WithSessionAdmission(ctx, sessionID, fn) +} + +func (q *failSecondCleanupUpsertQueue) RemoveEntryWithEntry( + ctx context.Context, + sessionID, entryID string, +) (*messagequeue.QueuedMessage, error) { + return q.service.RemoveEntryWithEntry(ctx, sessionID, entryID) +} + +func TestAdmissionClaimFailureRemovesEntryWhenCleanupStateUpdateFails(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + handlers, queue, db := newPersistentCleanupQueue(t, dbPath) + handlers.queueService = &failSecondCleanupUpsertQueue{QueueService: queue, service: queue} + claimer := &admissionCleanupAssertingClaimer{store: queue} + handlers.SetAttachmentClaimer(claimer) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + + response, err := handlers.wsQueueMessage(ctx, createTestMessage(t, ws.ActionMessageQueueAdd, map[string]interface{}{ + "session_id": "session-admission-update-failure", + "task_id": "task-admission-update-failure", + "content": "queued", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-admission-update-failure", + Name: "attachment.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + require.Empty(t, queue.GetStatus(ctx, "session-admission-update-failure").Entries) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + recoveryClaimer := &admissionClaimRecoveryClaimer{} + restarted.SetAttachmentClaimer(recoveryClaimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + cleanups, listErr := restartedQueue.ListAttachmentCleanups(ctx) + return listErr == nil && len(cleanups) == 0 && recoveryClaimer.releases.Load() == 1 + }, time.Second, 10*time.Millisecond) + require.Zero(t, recoveryClaimer.claims.Load()) +} + +func TestAdmissionCleanupDoesNotMutateReplacementEntryAfterRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + _, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := context.Background() + original, err := queue.QueueMessage( + ctx, "session-cleanup-fence", "task-cleanup-fence", "original", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-original", + Name: "original.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + fingerprint, err := queuedMessageFingerprint(original) + require.NoError(t, err) + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: original.SessionID, EntryID: original.ID, OperationID: "admission-fenced", + TaskID: original.TaskID, OwnerID: "owner", RemoveEntry: true, + EntryFingerprint: fingerprint, Attachments: original.Attachments, + })) + require.NoError(t, queue.UpdateMessageWithMetadata( + ctx, original.SessionID, original.ID, "replacement", + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-replacement", + Name: "replacement.txt", MimeType: "text/plain", SizeBytes: 1, + }}, nil, messagequeue.QueuedByUser, + )) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + claimer := &admissionClaimRecoveryClaimer{} + restarted.SetAttachmentClaimer(claimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + cleanups, listErr := restartedQueue.ListAttachmentCleanups(ctx) + return listErr == nil && len(cleanups) == 0 && claimer.releases.Load() == 1 + }, time.Second, 10*time.Millisecond) + current, err := restartedQueue.GetEntry(ctx, original.SessionID, original.ID) + require.NoError(t, err) + require.Equal(t, "replacement", current.Content) + require.Zero(t, claimer.claims.Load()) +} + +func TestAdmissionCleanupStillRemovesReorderedEntryAfterRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + _, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := context.Background() + target, err := queue.QueueMessage( + ctx, "session-cleanup-reorder", "task-cleanup-reorder", "target", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-reorder", + Name: "reorder.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + other, err := queue.QueueMessage( + ctx, target.SessionID, target.TaskID, "other", "", + messagequeue.QueuedByUser, false, nil, + ) + require.NoError(t, err) + fingerprint, err := queuedMessageFingerprint(target) + require.NoError(t, err) + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: target.SessionID, EntryID: target.ID, OperationID: "admission-reordered", + TaskID: target.TaskID, OwnerID: "owner", RemoveEntry: true, + EntryFingerprint: fingerprint, Attachments: target.Attachments, + })) + require.NoError(t, queue.ReorderEntries(ctx, target.SessionID, []string{other.ID, target.ID})) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + claimer := &admissionClaimRecoveryClaimer{} + restarted.SetAttachmentClaimer(claimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + _, getErr := restartedQueue.GetEntry(ctx, target.SessionID, target.ID) + return errors.Is(getErr, messagequeue.ErrEntryNotFound) && claimer.releases.Load() == 1 + }, time.Second, 10*time.Millisecond) +} + +func TestAdmissionCleanupFollowsTransferredEntryAfterRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + _, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := context.Background() + entry, err := queue.QueueMessage( + ctx, "session-cleanup-transfer-old", "task-cleanup-transfer", "queued", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-transfer", + Name: "transfer.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + fingerprint, err := queuedMessageFingerprint(entry) + require.NoError(t, err) + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "admission-transferred", + TaskID: entry.TaskID, OwnerID: "owner", RemoveEntry: true, + EntryFingerprint: fingerprint, Attachments: entry.Attachments, + })) + const destinationSessionID = "session-cleanup-transfer-new" + require.NoError(t, queue.TransferSession(ctx, entry.SessionID, destinationSessionID)) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + claimer := &admissionClaimRecoveryClaimer{} + restarted.SetAttachmentClaimer(claimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + _, getErr := restartedQueue.GetEntry(ctx, destinationSessionID, entry.ID) + cleanups, listErr := restartedQueue.ListAttachmentCleanups(ctx) + return errors.Is(getErr, messagequeue.ErrEntryNotFound) && + claimer.releases.Load() == 1 && listErr == nil && len(cleanups) == 0 + }, time.Second, 10*time.Millisecond) + cleanups, err := restartedQueue.ListAttachmentCleanups(ctx) + require.NoError(t, err) + require.Empty(t, cleanups) +} + +type transferReadObservingQueue struct { + *messagequeue.Service + destinationSessionID string + destinationReads atomic.Int32 +} + +func (q *transferReadObservingQueue) GetEntry( + ctx context.Context, + sessionID, entryID string, +) (*messagequeue.QueuedMessage, error) { + if sessionID == q.destinationSessionID { + q.destinationReads.Add(1) + } + return q.Service.GetEntry(ctx, sessionID, entryID) +} + +func TestTransferredAdmissionCleanupRetainsObligationOnContentMismatch(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + _, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := context.Background() + entry, err := queue.QueueMessage( + ctx, "session-cleanup-transfer-mismatch-old", "task-cleanup-transfer-mismatch", "original", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-transfer-mismatch", + Name: "transfer-mismatch.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + fingerprint, err := queuedMessageFingerprint(entry) + require.NoError(t, err) + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "admission-transfer-mismatch", + TaskID: entry.TaskID, OwnerID: "owner", RemoveEntry: true, + EntryFingerprint: fingerprint, Attachments: entry.Attachments, + })) + const destinationSessionID = "session-cleanup-transfer-mismatch-new" + require.NoError(t, queue.TransferSession(ctx, entry.SessionID, destinationSessionID)) + require.NoError(t, queue.UpdateMessageWithMetadata( + ctx, destinationSessionID, entry.ID, "replacement", entry.Attachments, nil, messagequeue.QueuedByUser, + )) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + observedQueue := &transferReadObservingQueue{ + Service: restartedQueue, destinationSessionID: destinationSessionID, + } + restarted.queueService = observedQueue + claimer := &admissionClaimRecoveryClaimer{} + restarted.SetAttachmentClaimer(claimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + return observedQueue.destinationReads.Load() >= 2 + }, time.Second, 10*time.Millisecond) + current, err := restartedQueue.GetEntry(ctx, destinationSessionID, entry.ID) + require.NoError(t, err) + require.Equal(t, "replacement", current.Content) + cleanups, err := restartedQueue.ListAttachmentCleanups(ctx) + require.NoError(t, err) + require.Len(t, cleanups, 1) + require.Zero(t, claimer.releases.Load()) +} + +func TestAdmissionCleanupWaitsForAnyActiveEditLease(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + ctx := context.Background() + entry, err := queue.QueueMessage( + ctx, "session-admission-lease", "task-admission-lease", "queued", "", + messagequeue.QueuedByUser, false, nil, + ) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + pending := &pendingQueueAttachmentCleanup{ + req: wsUpdateMessageRequest{ + SessionID: entry.SessionID, EntryID: entry.ID, LeaseID: "superseded-lease", + }, + authCtx: ctx, + } + + require.True(t, handlers.editLeaseActive(pending)) + require.NoError(t, queue.EndEdit(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection")) + require.False(t, handlers.editLeaseActive(pending)) +} + +type leaseAfterCleanupPrecheckQueue struct { + QueueService + service *messagequeue.Service + active atomic.Bool +} + +func (q *leaseAfterCleanupPrecheckQueue) GetEditLease( + context.Context, + string, + string, +) (*messagequeue.QueueEditLease, error) { + if !q.active.Load() { + return nil, messagequeue.ErrEditLeaseNotFound + } + return &messagequeue.QueueEditLease{LeaseID: "successor-lease"}, nil +} + +func (q *leaseAfterCleanupPrecheckQueue) WithSessionAdmission( + ctx context.Context, + sessionID string, + fn func(context.Context) error, +) error { + q.active.Store(true) + return q.service.WithSessionAdmission(ctx, sessionID, fn) +} +func TestAdmissionCleanupRechecksLeaseInsideAdmission(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + entry, err := queue.QueueMessage( + ctx, "session-cleanup-toctou", "task-cleanup-toctou", "queued", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-toctou", + Name: "toctou.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + fingerprint, err := queuedMessageFingerprint(entry) + require.NoError(t, err) + releaser := &controlledCleanupClaimer{} + wrapped := &leaseAfterCleanupPrecheckQueue{QueueService: queue, service: queue} + handlers.queueService = wrapped + pending := &pendingQueueAttachmentCleanup{ + req: wsUpdateMessageRequest{ + SessionID: entry.SessionID, EntryID: entry.ID, LeaseID: "superseded-lease", + }, + previous: &messagequeue.QueuedMessage{ + ID: entry.ID, SessionID: entry.SessionID, TaskID: entry.TaskID, + Attachments: entry.Attachments, + }, + releaser: releaser, removeEntry: true, claimPending: true, + entryFingerprint: fingerprint, authCtx: ctx, + } + + require.False(t, handlers.editLeaseActive(pending)) + require.Error(t, handlers.runPendingAttachmentCleanup(pending)) + _, err = queue.GetEntry(ctx, entry.SessionID, entry.ID) + require.NoError(t, err) + require.Zero(t, releaser.released.Load()) +} + +type insertionDuringCleanupReleaseClaimer struct { + queue *messagequeue.Service + sessionID string + taskID string + attachment messagequeue.MessageAttachment + insertedBeforeRelease atomic.Bool + insertDone chan error +} + +func (c *insertionDuringCleanupReleaseClaimer) ReleaseMessageAttachments( + context.Context, + string, + string, + []v1.MessageAttachment, +) error { + go func() { + _, err := c.queue.QueueMessage( + context.Background(), c.sessionID, c.taskID, "successor", "", + messagequeue.QueuedByUser, false, []messagequeue.MessageAttachment{c.attachment}, + ) + c.insertDone <- err + }() + select { + case err := <-c.insertDone: + c.insertedBeforeRelease.Store(true) + return err + case <-time.After(50 * time.Millisecond): + return nil + } +} + +func TestMissingEntryCleanupHoldsAdmissionThroughRelease(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + attachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "attachment-missing-admission", + Name: "missing-admission.txt", MimeType: "text/plain", SizeBytes: 1, + } + releaser := &insertionDuringCleanupReleaseClaimer{ + queue: queue, sessionID: "session-missing-admission", taskID: "task-missing-admission", + attachment: attachment, insertDone: make(chan error, 1), + } + pending := &pendingQueueAttachmentCleanup{ + key: pendingQueueAttachmentCleanupKey{ + sessionID: releaser.sessionID, entryID: "missing-entry", operationID: "missing-admission", + }, + req: wsUpdateMessageRequest{ + SessionID: releaser.sessionID, EntryID: "missing-entry", OperationID: "missing-admission", + }, + previous: &messagequeue.QueuedMessage{ + ID: "missing-entry", SessionID: releaser.sessionID, TaskID: releaser.taskID, + Attachments: []messagequeue.MessageAttachment{attachment}, + }, + releaser: releaser, removeEntry: true, claimPending: true, + authCtx: authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}), + } + + require.NoError(t, handlers.runPendingAttachmentCleanup(pending)) + require.False(t, releaser.insertedBeforeRelease.Load()) + select { + case err := <-releaser.insertDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("successor insertion remained blocked after cleanup release") + } +} + +func TestTransferredSupersededCleanupChecksDestinationLease(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + oldAttachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "attachment-transfer-old", + Name: "old.txt", MimeType: "text/plain", SizeBytes: 1, + } + newAttachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "attachment-transfer-new", + Name: "new.txt", MimeType: "text/plain", SizeBytes: 1, + } + entry, err := queue.QueueMessage( + ctx, "session-superseded-transfer-old", "task-superseded-transfer", "updated", "", + messagequeue.QueuedByUser, false, []messagequeue.MessageAttachment{newAttachment}, + ) + require.NoError(t, err) + const destinationSessionID = "session-superseded-transfer-new" + require.NoError(t, queue.TransferSession(ctx, entry.SessionID, destinationSessionID)) + lease, err := queue.BeginEdit(ctx, destinationSessionID, entry.ID, "successor-connection") + require.NoError(t, err) + releaser := &controlledCleanupClaimer{} + pending := &pendingQueueAttachmentCleanup{ + key: pendingQueueAttachmentCleanupKey{ + sessionID: entry.SessionID, entryID: entry.ID, operationID: "superseded-transfer", + }, + req: wsUpdateMessageRequest{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "superseded-transfer", + }, + previous: &messagequeue.QueuedMessage{ + ID: entry.ID, SessionID: entry.SessionID, TaskID: entry.TaskID, + Attachments: []messagequeue.MessageAttachment{oldAttachment}, + }, + releaser: releaser, authCtx: ctx, + } + + require.Error(t, handlers.runPendingAttachmentCleanup(pending)) + require.Equal(t, destinationSessionID, handlers.pendingAttachmentCleanupSessionID(pending)) + require.Zero(t, releaser.released.Load()) + require.NoError(t, queue.EndEdit(ctx, destinationSessionID, entry.ID, lease.LeaseID, "successor-connection")) +} + +func TestTransferredCleanupWakesWhenDestinationEditEnds(t *testing.T) { + handlers, queue, db := newPersistentCleanupQueue(t, filepath.Join(t.TempDir(), "queue.db")) + defer func() { _ = db.Close() }() + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + attachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "attachment-wake-transfer", + Name: "wake-transfer.txt", MimeType: "text/plain", SizeBytes: 1, + } + entry, err := queue.QueueMessage( + ctx, "session-wake-transfer-old", "task-wake-transfer", "queued", "", + messagequeue.QueuedByUser, false, []messagequeue.MessageAttachment{attachment}, + ) + require.NoError(t, err) + key := pendingQueueAttachmentCleanupKey{ + sessionID: entry.SessionID, entryID: entry.ID, operationID: "wake-transfer", + } + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: key.operationID, + TaskID: entry.TaskID, OwnerID: "owner", Attachments: []messagequeue.MessageAttachment{attachment}, + })) + pending := &pendingQueueAttachmentCleanup{ + key: key, + req: wsUpdateMessageRequest{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: key.operationID, + }, + currentSessionID: entry.SessionID, + authCtx: ctx, + wake: make(chan struct{}, 1), + } + handlers.attachmentCleanupMu.Lock() + handlers.pendingAttachmentCleanup[key] = pending + handlers.attachmentCleanupMu.Unlock() + const destinationSessionID = "session-wake-transfer-new" + require.NoError(t, queue.TransferSession(ctx, entry.SessionID, destinationSessionID)) + const connectionID = "connection-wake-transfer" + lease, err := queue.BeginEdit(ctx, destinationSessionID, entry.ID, connectionID) + require.NoError(t, err) + + response, err := handlers.wsEndEdit( + ws.WithConnectionID(ctx, connectionID), + createTestMessage(t, ws.ActionMessageQueueEditEnd, map[string]string{ + "session_id": destinationSessionID, + "entry_id": entry.ID, + "lease_id": lease.LeaseID, + }), + ) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + select { + case <-pending.wake: + case <-time.After(time.Second): + t.Fatal("destination edit end did not wake transferred cleanup") + } +} + +func TestMissingEntryCleanupUsesTransferredClaimSessionAfterRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + handlers, queue, db := newPersistentCleanupQueue(t, dbPath) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + const ( + sourceSessionID = "session-missing-transfer-old" + destinationSessionID = "session-missing-transfer-new" + taskID = "task-missing-transfer" + entryID = "missing-transfer-entry" + operationID = "missing-transfer-cleanup" + ) + attachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "attachment-missing-transfer", + Name: "missing-transfer.txt", MimeType: "text/plain", SizeBytes: 1, + } + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: sourceSessionID, EntryID: entryID, OperationID: operationID, + TaskID: taskID, OwnerID: "owner", RemoveEntry: true, ClaimPending: true, + Attachments: []messagequeue.MessageAttachment{attachment}, + })) + firstClaimer := &controlledCleanupClaimer{} + firstClaimer.failures.Store(1) + pending := &pendingQueueAttachmentCleanup{ + key: pendingQueueAttachmentCleanupKey{ + sessionID: sourceSessionID, entryID: entryID, operationID: operationID, + }, + req: wsUpdateMessageRequest{ + SessionID: sourceSessionID, EntryID: entryID, OperationID: operationID, + }, + previous: &messagequeue.QueuedMessage{ + ID: entryID, SessionID: sourceSessionID, TaskID: taskID, + Attachments: []messagequeue.MessageAttachment{attachment}, + }, + releaser: firstClaimer, removeEntry: true, claimPending: true, authCtx: ctx, + } + require.Error(t, handlers.runPendingAttachmentCleanup(pending)) + require.NoError(t, queue.TransferSession(ctx, sourceSessionID, destinationSessionID)) + require.NoError(t, db.Close()) + + restarted, restartedQueue, restartedDB := newPersistentCleanupQueue(t, dbPath) + defer func() { _ = restartedDB.Close() }() + recoveryClaimer := &controlledCleanupClaimer{} + restarted.SetAttachmentClaimer(recoveryClaimer) + restarted.Start(context.Background()) + defer restarted.Stop() + + require.Eventually(t, func() bool { + cleanups, err := restartedQueue.ListAttachmentCleanups(ctx) + return err == nil && len(cleanups) == 0 && recoveryClaimer.attempts.Load() > 0 + }, time.Second, 10*time.Millisecond) + require.Equal(t, destinationSessionID, recoveryClaimer.lastSession.Load()) +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_cleanup_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_cleanup_test.go index 9d9eeeea2e..3f30206768 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_cleanup_test.go +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_cleanup_test.go @@ -3,17 +3,758 @@ package handlers import ( "context" "database/sql" + "encoding/json" + "errors" "path/filepath" + "sync/atomic" "testing" + "time" "github.com/jmoiron/sqlx" - "github.com/kandev/kandev/internal/common/logger" - "github.com/kandev/kandev/internal/orchestrator/messagequeue" - ws "github.com/kandev/kandev/pkg/websocket" _ "github.com/mattn/go-sqlite3" + + "github.com/kandev/kandev/internal/auth/authn" + "github.com/kandev/kandev/internal/common/logger" + eventtypes "github.com/kandev/kandev/internal/events" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + v1 "github.com/kandev/kandev/pkg/api/v1" + ws "github.com/kandev/kandev/pkg/websocket" "github.com/stretchr/testify/require" ) +func TestWsRemoveEntryReleasesClaimedAttachments(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &recordingQueueAttachmentClaimer{} + handlers.SetAttachmentClaimer(claimer) + entry, err := queue.QueueMessage(context.Background(), "session", "task", "queued", "", "user", false, []messagequeue.MessageAttachment{{ + Type: "resource", + AttachmentID: "attachment", + Name: "report.txt", + MimeType: "text/plain", + }}) + require.NoError(t, err) + + response, err := handlers.wsRemoveEntry(context.Background(), createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": "session", + "entry_id": entry.ID, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Equal(t, []string{"attachment"}, claimer.releases) +} + +func TestWsRemoveEntryPreservesAttachmentReferencedByAnotherEntry(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &recordingQueueAttachmentClaimer{} + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + attachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "shared-attachment", Name: "shared.txt", MimeType: "text/plain", + } + first, err := queue.QueueMessage(ctx, "session-shared", "task-shared", "first", "", "user", false, []messagequeue.MessageAttachment{attachment}) + require.NoError(t, err) + _, err = queue.QueueMessage(ctx, "session-shared", "task-shared", "second", "", "user", false, []messagequeue.MessageAttachment{attachment}) + require.NoError(t, err) + + response, err := handlers.wsRemoveEntry(ctx, createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": "session-shared", + "entry_id": first.ID, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Empty(t, claimer.releases) +} + +func TestWsRemoveEntryUsesAtomicRemovedEntryForAttachmentCleanup(t *testing.T) { + handlers, _ := setupQueueHandlers(t) + stale := &messagequeue.QueuedMessage{ + ID: "entry", SessionID: "session", TaskID: "task", + Attachments: []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "stale", Name: "stale.txt", MimeType: "text/plain", + }}, + } + current := *stale + current.Attachments = []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "current", Name: "current.txt", MimeType: "text/plain", + }} + queue := &atomicRemoveQueueService{ + QueueService: handlers.queueService, + stale: stale, + removed: ¤t, + } + handlers.queueService = queue + claimer := &recordingQueueAttachmentClaimer{} + handlers.SetAttachmentClaimer(claimer) + + response, err := handlers.wsRemoveEntry(context.Background(), createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": "session", + "entry_id": "entry", + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Equal(t, 1, queue.atomicCalls) + require.Equal(t, []string{"current"}, claimer.releases) +} + +type atomicRemoveQueueService struct { + QueueService + stale, removed *messagequeue.QueuedMessage + atomicCalls int +} + +func (s *atomicRemoveQueueService) GetEntry(context.Context, string, string) (*messagequeue.QueuedMessage, error) { + return s.stale, nil +} + +func (s *atomicRemoveQueueService) RemoveEntry(context.Context, string, string) error { + return nil +} + +func (s *atomicRemoveQueueService) RemoveEntryWithEntry(context.Context, string, string) (*messagequeue.QueuedMessage, error) { + s.atomicCalls++ + return s.removed, nil +} + +func TestWsCancelAllReleasesClaimedAttachments(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &recordingQueueAttachmentClaimer{} + handlers.SetAttachmentClaimer(claimer) + for _, attachmentID := range []string{"first", "second"} { + _, err := queue.QueueMessage(context.Background(), "session-cancel", "task-cancel", "queued", "", "user", false, []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: attachmentID, Name: attachmentID + ".txt", MimeType: "text/plain", + }}) + require.NoError(t, err) + } + + response, err := handlers.wsCancelAll(context.Background(), createTestMessage(t, ws.ActionMessageQueueCancel, map[string]string{ + "session_id": "session-cancel", + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.ElementsMatch(t, []string{"first", "second"}, claimer.releases) +} +func TestWsCancelAllReportsOnlyRemovedEntries(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + ctx := context.Background() + + _, _, accepted, err := queue.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-cancel-count", "task", "lifecycle", "", "", false, nil, + nil, "lifecycle", true, + ) + require.NoError(t, err) + require.True(t, accepted) + _, ok := queue.ReserveQueued(ctx, "session-cancel-count") + require.True(t, ok) + _, err = queue.QueueMessage(ctx, "session-cancel-count", "task", "ordinary", "", "user", false, nil) + require.NoError(t, err) + + response, err := handlers.wsCancelAll(ctx, createTestMessage(t, ws.ActionMessageQueueCancel, map[string]string{ + "session_id": "session-cancel-count", + })) + require.NoError(t, err) + + var payload map[string]interface{} + require.NoError(t, json.Unmarshal(response.Payload, &payload)) + require.EqualValues(t, 1, payload["removed"]) + require.Equal(t, 0, queue.GetStatus(ctx, "session-cancel-count").Count) +} + +func TestWsUpdateMessageRollsBackBeforeSuccessorCanAcquireEdit(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session", "task", "original", "", "user", false, []messagequeue.MessageAttachment{{ + Type: "resource", + AttachmentID: "new-attachment", + Name: "new.txt", + MimeType: "text/plain", + }}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + claimer := &successorEditRaceClaimer{ + queue: queue, + sessionID: entry.SessionID, + entryID: entry.ID, + leaseID: lease.LeaseID, + successorID: "connection-b", + successorDone: make(chan struct{}), + } + handlers.SetAttachmentClaimer(claimer) + response, err := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-a"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": "session", + "entry_id": entry.ID, + "lease_id": lease.LeaseID, + "operation_id": "operation-1", + "expected_target_revision": lease.TargetRevision, + "content": "edited", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new-attachment", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + <-claimer.successorDone + require.False(t, claimer.successBeforeRelease.Load()) + require.True(t, claimer.releaseObserved.Load()) +} +func TestWsUpdateMessageWithLeaseReleasesSupersededAttachments(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &recordingQueueAttachmentClaimer{} + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session-update", "task-update", "before", "", "user", false, []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "old-attachment", Name: "old.txt", MimeType: "text/plain", + }}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-a"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + "lease_id": lease.LeaseID, + "operation_id": "operation-update", + "expected_target_revision": lease.TargetRevision, + "content": "after", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new-attachment", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Equal(t, []string{"new-attachment"}, claimer.claims) + require.Equal(t, []string{"old-attachment"}, claimer.releases) +} + +type failingThenSuccessfulQueueAttachmentClaimer struct { + failures atomic.Int32 + attempts atomic.Int32 + released atomic.Int32 +} + +func (c *failingThenSuccessfulQueueAttachmentClaimer) ClaimMessageAttachments(context.Context, string, string, []v1.MessageAttachment) error { + return nil +} + +func (c *failingThenSuccessfulQueueAttachmentClaimer) ReleaseMessageAttachments(_ context.Context, _ string, _ string, attachments []v1.MessageAttachment) error { + c.attempts.Add(1) + if c.failures.Load() > 0 { + c.failures.Add(-1) + return errors.New("transient attachment release failure") + } + for _, attachment := range attachments { + if attachment.AttachmentID != "" { + c.released.Add(1) + } + } + return nil +} + +func TestWsUpdateMessageRetriesSupersededAttachmentCleanupAfterLeaseEnd(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &failingThenSuccessfulQueueAttachmentClaimer{} + claimer.failures.Store(1) + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session-retry-cleanup", "task-retry-cleanup", "before", "", "user", false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "old-attachment", Name: "old.txt", MimeType: "text/plain", + }}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-retry") + require.NoError(t, err) + message := func() *ws.Message { + return createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + "lease_id": lease.LeaseID, + "operation_id": "operation-retry", + "expected_target_revision": lease.TargetRevision, + "content": "after", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new-attachment", Name: "new.txt", MimeType: "text/plain", + }}, + }) + } + + firstResponse, err := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-retry"), message()) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, firstResponse.Type) + require.NoError(t, queue.EndEdit(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection-retry")) + require.Eventually(t, func() bool { + return claimer.attempts.Load() >= 2 && claimer.released.Load() == 1 + }, time.Second, 10*time.Millisecond) + require.Equal(t, int32(2), claimer.attempts.Load()) + require.Equal(t, int32(1), claimer.released.Load()) +} + +func TestWsRemoveEntryRetriesFailedAttachmentRelease(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &failingThenSuccessfulQueueAttachmentClaimer{} + claimer.failures.Store(1) + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session-remove-retry", "task-remove-retry", "queued", "", "user", false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment-remove"}}) + require.NoError(t, err) + + response, err := handlers.wsRemoveEntry(ctx, createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Eventually(t, func() bool { + return claimer.attempts.Load() == 2 && claimer.released.Load() == 1 + }, time.Second, 10*time.Millisecond) +} + +func TestWsCancelAllRetriesFailedAttachmentRelease(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &failingThenSuccessfulQueueAttachmentClaimer{} + claimer.failures.Store(1) + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + _, err := queue.QueueMessage(ctx, "session-cancel-retry", "task-cancel-retry", "queued", "", "user", false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment-cancel"}}) + require.NoError(t, err) + + response, err := handlers.wsCancelAll(ctx, createTestMessage(t, ws.ActionMessageQueueCancel, map[string]string{ + "session_id": "session-cancel-retry", + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Eventually(t, func() bool { + return claimer.attempts.Load() == 2 && claimer.released.Load() == 1 + }, time.Second, 10*time.Millisecond) +} + +type failingReferenceQueueService struct { + QueueService + service *messagequeue.Service + attempts atomic.Int32 +} + +func (s *failingReferenceQueueService) ReferencedQueueAttachmentIDs( + ctx context.Context, + sessionID, excludedEntryID string, + attachmentIDs []string, +) (map[string]struct{}, error) { + if s.attempts.Add(1) == 1 { + return nil, errors.New("reference lookup unavailable") + } + return s.service.ReferencedQueueAttachmentIDs(ctx, sessionID, excludedEntryID, attachmentIDs) +} + +func (s *failingReferenceQueueService) RemoveEntryWithEntry( + ctx context.Context, + sessionID, entryID string, +) (*messagequeue.QueuedMessage, error) { + return s.service.RemoveEntryWithEntry(ctx, sessionID, entryID) +} + +func TestWsRemoveEntryRetriesFailedAttachmentReferenceLookup(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + references := &failingReferenceQueueService{QueueService: queue, service: queue} + handlers.queueService = references + claimer := &controlledCleanupClaimer{} + handlers.SetAttachmentClaimer(claimer) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-reference-retry"}) + handlers.Start(ctx) + t.Cleanup(handlers.Stop) + entry, err := queue.QueueMessage( + ctx, "session-reference-retry", "task-reference-retry", "queued", "", "user", false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment-reference-retry"}}, + ) + require.NoError(t, err) + + response, err := handlers.wsRemoveEntry( + ctx, + createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + }), + ) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Eventually(t, func() bool { + return references.attempts.Load() >= 2 && claimer.released.Load() == 1 + }, time.Second, 10*time.Millisecond) +} + +func TestPendingAttachmentCleanupResumesAfterHandlerRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + handlers, queue, db := newPersistentCleanupQueue(t, dbPath) + firstClaimer := &controlledCleanupClaimer{} + firstClaimer.failures.Store(100) + handlers.SetAttachmentClaimer(firstClaimer) + handlers.Start(context.Background()) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-restart"}) + entry, err := queue.QueueMessage(ctx, "session-cleanup-restart", "task-cleanup-restart", "before", "", "user", false, + []messagequeue.MessageAttachment{{AttachmentID: "old-attachment"}}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-restart") + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-restart"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, "entry_id": entry.ID, "lease_id": lease.LeaseID, + "operation_id": "operation-restart", "expected_target_revision": lease.TargetRevision, + "content": "after", "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new-attachment", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Equal(t, int32(1), firstClaimer.attempts.Load()) + handlers.Stop() + require.NoError(t, db.Close()) + + restarted, _, restartedDB := newPersistentCleanupQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + secondClaimer := &controlledCleanupClaimer{} + restarted.SetAttachmentClaimer(secondClaimer) + restarted.Start(context.Background()) + t.Cleanup(restarted.Stop) + + require.Eventually(t, func() bool { return secondClaimer.released.Load() == 1 }, time.Second, 10*time.Millisecond) + require.Equal(t, "user-restart", secondClaimer.lastUser.Load()) +} +func TestLegacyAttachmentCleanupWithoutOwnerIsSkippedWithoutInternalReleaser(t *testing.T) { + handlers, queue, db := newPersistentCleanupQueue(t, filepath.Join(t.TempDir(), "queue.db")) + t.Cleanup(func() { _ = db.Close() }) + ctx := context.Background() + entry, err := queue.QueueMessage( + ctx, "session-legacy-cleanup", "task-legacy-cleanup", "queued", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "legacy-attachment"}}, + ) + require.NoError(t, err) + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "legacy-cleanup", + TaskID: entry.TaskID, Attachments: entry.Attachments, + })) + + handlers.SetAttachmentClaimer(&controlledCleanupClaimer{}) + handlers.Start(ctx) + t.Cleanup(handlers.Stop) + + handlers.attachmentCleanupMu.Lock() + pending := len(handlers.pendingAttachmentCleanup) + handlers.attachmentCleanupMu.Unlock() + require.Zero(t, pending) +} + +type cleanupPersistenceFailureQueue struct { + QueueService + service *messagequeue.Service + err error +} + +func (q *cleanupPersistenceFailureQueue) AttachmentCleanupPersistenceAvailable() bool { return true } +func (q *cleanupPersistenceFailureQueue) UpsertAttachmentCleanup(context.Context, messagequeue.AttachmentCleanup) error { + return q.err +} +func (q *cleanupPersistenceFailureQueue) DeleteAttachmentCleanup(context.Context, string, string, string) error { + return nil +} +func (q *cleanupPersistenceFailureQueue) ListAttachmentCleanups(context.Context) ([]messagequeue.AttachmentCleanup, error) { + return nil, nil +} +func (q *cleanupPersistenceFailureQueue) WithSessionAdmission( + ctx context.Context, + sessionID string, + fn func(context.Context) error, +) error { + return q.service.WithSessionAdmission(ctx, sessionID, fn) +} +func (q *cleanupPersistenceFailureQueue) RemoveEntryWithEntry( + ctx context.Context, + sessionID, entryID string, +) (*messagequeue.QueuedMessage, error) { + return q.service.RemoveEntryWithEntry(ctx, sessionID, entryID) +} + +func TestWsRemoveEntryPreservesQueueWhenCleanupPersistenceFails(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + handlers.queueService = &cleanupPersistenceFailureQueue{ + QueueService: queue, + service: queue, + err: errors.New("cleanup store unavailable"), + } + handlers.SetAttachmentClaimer(failingQueueAttachmentReleaser{}) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-persist"}) + entry, err := queue.QueueMessage(ctx, "session-persist-failure", "task-persist-failure", "queued", "", "user", false, + []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}) + require.NoError(t, err) + + response, err := handlers.wsRemoveEntry(ctx, createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + _, err = queue.GetEntry(ctx, entry.SessionID, entry.ID) + require.NoError(t, err) +} + +func TestFailedEditRollbackCleanupResumesAfterRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + handlers, queue, db := newPersistentCleanupQueue(t, dbPath) + firstClaimer := &controlledCleanupClaimer{} + firstClaimer.failures.Store(100) + handlers.SetAttachmentClaimer(firstClaimer) + handlers.Start(context.Background()) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-rollback"}) + entry, err := queue.QueueMessage(ctx, "session-rollback-restart", "task-rollback-restart", "before", "", "user", false, + []messagequeue.MessageAttachment{{AttachmentID: "old-attachment"}}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-rollback") + require.NoError(t, err) + _, err = db.Exec(` + CREATE TRIGGER fail_queue_attachment_update + BEFORE UPDATE ON queued_messages + BEGIN + SELECT RAISE(ABORT, 'forced queue update failure'); + END + `) + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-rollback"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, "entry_id": entry.ID, "lease_id": lease.LeaseID, + "operation_id": "operation-rollback", "expected_target_revision": lease.TargetRevision, + "content": "after", "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new-attachment", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + require.Equal(t, int32(1), firstClaimer.attempts.Load()) + handlers.Stop() + require.NoError(t, db.Close()) + + restarted, _, restartedDB := newPersistentCleanupQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + secondClaimer := &controlledCleanupClaimer{} + restarted.SetAttachmentClaimer(secondClaimer) + restarted.Start(context.Background()) + t.Cleanup(restarted.Stop) + + require.Eventually(t, func() bool { return secondClaimer.released.Load() == 1 }, time.Second, 10*time.Millisecond) + require.Equal(t, "user-rollback", secondClaimer.lastUser.Load()) +} + +func TestLeaseClaimFailureCleanupResumesAfterRestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "queue.db") + handlers, queue, db := newPersistentCleanupQueue(t, dbPath) + firstClaimer := &controlledCleanupClaimer{claimErr: errors.New("claim failed")} + firstClaimer.failures.Store(100) + handlers.SetAttachmentClaimer(firstClaimer) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-claim-rollback"}) + entry, err := queue.QueueMessage( + ctx, "session-claim-rollback", "task-claim-rollback", "before", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "old-attachment"}}, + ) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-claim-rollback") + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage( + ws.WithConnectionID(ctx, "connection-claim-rollback"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, "entry_id": entry.ID, "lease_id": lease.LeaseID, + "operation_id": "operation-claim-rollback", "expected_target_revision": lease.TargetRevision, + "content": "after", "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new-attachment", Name: "new.txt", MimeType: "text/plain", + }}, + }), + ) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + cleanups, err := queue.ListAttachmentCleanups(ctx) + require.NoError(t, err) + require.Len(t, cleanups, 1, "failed release must retain the durable rollback obligation") + require.NoError(t, db.Close()) + + restarted, _, restartedDB := newPersistentCleanupQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + secondClaimer := &controlledCleanupClaimer{} + restarted.SetAttachmentClaimer(secondClaimer) + restarted.Start(context.Background()) + t.Cleanup(restarted.Stop) + + require.Eventually(t, func() bool { return secondClaimer.released.Load() == 1 }, time.Second, 10*time.Millisecond) + require.Equal(t, "user-claim-rollback", secondClaimer.lastUser.Load()) +} + +func newPersistentCleanupQueue(t *testing.T, dbPath string) (*QueueHandlers, *messagequeue.Service, *sqlx.DB) { + t.Helper() + raw, err := sql.Open("sqlite3", dbPath) + require.NoError(t, err) + raw.SetMaxOpenConns(1) + db := sqlx.NewDb(raw, "sqlite3") + repo, err := messagequeue.NewSQLiteRepository(db, db) + require.NoError(t, err) + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console", OutputPath: "stderr"}) + require.NoError(t, err) + queue := messagequeue.NewService(repo, messagequeue.DefaultMaxPerSession, log) + queue.SetAutoMergeEnabled(false) + handlers := NewQueueHandlers(queue, &mockEventBus{}, log, nil, allowQueueAccess{}, nil) + return handlers, queue, db +} + +type controlledCleanupClaimer struct { + failures atomic.Int32 + attempts atomic.Int32 + released atomic.Int32 + lastSession atomic.Value + lastUser atomic.Value + claimErr error +} + +func (c *controlledCleanupClaimer) ClaimMessageAttachments(context.Context, string, string, []v1.MessageAttachment) error { + return c.claimErr +} + +func (c *controlledCleanupClaimer) ReleaseMessageAttachments(ctx context.Context, _ string, sessionID string, attachments []v1.MessageAttachment) error { + c.attempts.Add(1) + identity, ok := authn.IdentityFromContext(ctx) + if !ok || identity.UserID == "" { + return errors.New("missing attachment authorization") + } + if c.failures.Load() > 0 { + c.failures.Add(-1) + return errors.New("transient attachment release failure") + } + c.lastSession.Store(sessionID) + c.lastUser.Store(identity.UserID) + c.released.Add(int32(len(attachments))) + return nil +} + +func TestWsUpdateMessageCleanupRetryPreservesAuthAndPublishesStatus(t *testing.T) { + handlers, queue, events := setupQueueHandlersWithResolver(t, func(context.Context, string) (string, error) { + return "task-cleanup", nil + }) + claimer := &controlledCleanupClaimer{} + claimer.failures.Store(1) + handlers.SetAttachmentClaimer(claimer) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-cleanup"}) + entry, err := queue.QueueMessage(ctx, "session-cleanup", "task-cleanup", "before", "", "user", false, + []messagequeue.MessageAttachment{{Type: "resource", AttachmentID: "old", Name: "old.txt", MimeType: "text/plain"}}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-cleanup") + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-cleanup"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, "entry_id": entry.ID, "lease_id": lease.LeaseID, + "operation_id": "operation-cleanup", "expected_target_revision": lease.TargetRevision, + "content": "after", "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Equal(t, "session-cleanup", events.lastData[fieldSessionID]) + require.NoError(t, queue.EndEdit(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection-cleanup")) + require.Equal(t, eventtypes.MessageQueueStatusChanged, events.lastSubject) + require.Eventually(t, func() bool { return claimer.released.Load() == 1 }, time.Second, 10*time.Millisecond) + require.Equal(t, "user-cleanup", claimer.lastUser.Load()) +} + +func TestWsUpdateMessageCleanupRetryReconcilesRemovedEntry(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &controlledCleanupClaimer{} + claimer.failures.Store(1) + handlers.SetAttachmentClaimer(claimer) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-cleanup"}) + entry, err := queue.QueueMessage(ctx, "session-removed", "task-removed", "before", "", "user", false, + []messagequeue.MessageAttachment{{Type: "resource", AttachmentID: "old", Name: "old.txt", MimeType: "text/plain"}}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-removed") + require.NoError(t, err) + _, err = handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-removed"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, "entry_id": entry.ID, "lease_id": lease.LeaseID, + "operation_id": "operation-removed", "expected_target_revision": lease.TargetRevision, + "content": "after", "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.NoError(t, queue.RemoveEntry(ctx, entry.SessionID, entry.ID)) + require.Eventually(t, func() bool { return claimer.released.Load() == 1 }, time.Second, 10*time.Millisecond) +} + +func TestWsUpdateMessageCleanupRetryFollowsTransferredEntry(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &controlledCleanupClaimer{} + claimer.failures.Store(1) + handlers.SetAttachmentClaimer(claimer) + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "user-cleanup"}) + entry, err := queue.QueueMessage(ctx, "session-transfer-old", "task-transfer", "before", "", "user", false, + []messagequeue.MessageAttachment{{Type: "resource", AttachmentID: "old", Name: "old.txt", MimeType: "text/plain"}}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-transfer") + require.NoError(t, err) + _, err = handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-transfer"), + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, "entry_id": entry.ID, "lease_id": lease.LeaseID, + "operation_id": "operation-transfer", "expected_target_revision": lease.TargetRevision, + "content": "after", "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.NoError(t, queue.TransferSession(ctx, entry.SessionID, "session-transfer-new")) + require.Eventually(t, func() bool { return claimer.released.Load() == 1 }, time.Second, 10*time.Millisecond) + require.Equal(t, "session-transfer-new", claimer.lastSession.Load()) +} + +type successorEditRaceClaimer struct { + queue *messagequeue.Service + sessionID string + entryID string + leaseID string + successorID string + successorDone chan struct{} + successBeforeRelease atomic.Bool + releaseObserved atomic.Bool +} + +func (c *successorEditRaceClaimer) ClaimMessageAttachments(ctx context.Context, _, _ string, _ []v1.MessageAttachment) error { + if err := c.queue.EndEdit(ctx, c.sessionID, c.entryID, c.leaseID, "connection-a"); err != nil { + return err + } + go func() { + defer close(c.successorDone) + lease, err := c.queue.BeginEdit(context.Background(), c.sessionID, c.entryID, c.successorID) + if err != nil { + return + } + if !c.releaseObserved.Load() { + c.successBeforeRelease.Store(true) + } + _ = c.queue.EndEdit(context.Background(), c.sessionID, c.entryID, lease.LeaseID, c.successorID) + }() + select { + case <-c.successorDone: + case <-time.After(100 * time.Millisecond): + } + return nil +} + +func (c *successorEditRaceClaimer) ReleaseMessageAttachments(context.Context, string, string, []v1.MessageAttachment) error { + c.releaseObserved.Store(true) + return nil +} + func setupSQLiteQueueAttachmentCleanup(t *testing.T) (*QueueHandlers, *messagequeue.Service, messagequeue.Repository, messagequeue.QueueSessionIdentity, *recordingQueueAttachmentClaimer) { t.Helper() raw, err := sql.Open("sqlite3", filepath.Join(t.TempDir(), "queue.db")+"?_foreign_keys=on") diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_context_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_context_test.go new file mode 100644 index 0000000000..7d620a145e --- /dev/null +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_context_test.go @@ -0,0 +1,203 @@ +package handlers + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + v1 "github.com/kandev/kandev/pkg/api/v1" + ws "github.com/kandev/kandev/pkg/websocket" + "github.com/stretchr/testify/require" +) + +type contextRecordingReleaser struct { + err error + calls int +} + +func (r *contextRecordingReleaser) ClaimMessageAttachments(context.Context, string, string, []v1.MessageAttachment) error { + return nil +} + +func (r *contextRecordingReleaser) ReleaseMessageAttachments(ctx context.Context, _ string, _ string, _ []v1.MessageAttachment) error { + r.calls++ + r.err = ctx.Err() + return nil +} + +func TestQueueAttachmentCleanupDetachesCancelledContext(t *testing.T) { + handlers, _ := setupQueueHandlers(t) + releaser := &contextRecordingReleaser{} + handlers.SetAttachmentClaimer(releaser) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + entry := &messagequeue.QueuedMessage{ + ID: "entry", + SessionID: "session", + TaskID: "task", + Attachments: []messagequeue.MessageAttachment{{ + AttachmentID: "attachment", + }}, + } + pending, err := handlers.prepareEntryRemovalCleanup(ctx, entry, "remove") + require.NoError(t, err) + handlers.releaseQueuedAttachments(ctx, entry, pending) + + if releaser.err != nil { + t.Fatalf("attachment cleanup context error = %v, want nil", releaser.err) + } +} + +func TestWsRemoveEntryDetachesCancelledCleanupContext(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + entry, err := queue.QueueMessage(context.Background(), "session", "task", "queued", "", "user", false, []messagequeue.MessageAttachment{{ + AttachmentID: "attachment", + }}) + require.NoError(t, err) + cleanupQueue := &cleanupContextQueueService{QueueService: queue, entry: entry} + handlers.queueService = cleanupQueue + claimer := &contextRecordingReleaser{} + handlers.SetAttachmentClaimer(claimer) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + response, err := handlers.wsRemoveEntry(ctx, createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.NoError(t, cleanupQueue.referenceErr) + require.NoError(t, claimer.err) +} + +type cleanupContextQueueService struct { + QueueService + entry *messagequeue.QueuedMessage + getEntryErr error + referenceErr error +} + +func (s *cleanupContextQueueService) GetEntry(ctx context.Context, _, _ string) (*messagequeue.QueuedMessage, error) { + s.getEntryErr = ctx.Err() + return s.entry, nil +} + +func (s *cleanupContextQueueService) ReferencedQueueAttachmentIDs(ctx context.Context, _, _ string, _ []string) (map[string]struct{}, error) { + s.referenceErr = ctx.Err() + return nil, nil +} +func (s *cleanupContextQueueService) RemoveEntryWithEntry(context.Context, string, string) (*messagequeue.QueuedMessage, error) { + return s.entry, nil +} + +func TestQueueAttachmentCleanupDetachesReadContext(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + current := &messagequeue.QueuedMessage{ + ID: "entry", SessionID: "session", TaskID: "task", + Attachments: []messagequeue.MessageAttachment{{ + AttachmentID: "replacement", + }}, + } + cleanupQueue := &cleanupContextQueueService{QueueService: queue, entry: current} + handlers.queueService = cleanupQueue + releaser := &contextRecordingReleaser{} + previous := &messagequeue.QueuedMessage{ + ID: "entry", SessionID: "session", TaskID: "task", + Attachments: []messagequeue.MessageAttachment{{ + AttachmentID: "original", + }}, + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := handlers.releaseSupersededQueueAttachmentsAdmitted(ctx, wsUpdateMessageRequest{ + SessionID: "session", EntryID: "entry", + }, previous, releaser) + require.NoError(t, err) + + require.NoError(t, cleanupQueue.getEntryErr) + require.NoError(t, cleanupQueue.referenceErr) + require.NoError(t, releaser.err) +} + +func TestQueueAttachmentCleanupDetachesCancelledReferenceLookup(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + cleanupQueue := &cleanupContextQueueService{QueueService: queue} + handlers.queueService = cleanupQueue + releaser := &contextRecordingReleaser{} + previous := []messagequeue.MessageAttachment{{AttachmentID: "attachment"}} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := handlers.releaseQueueAttachmentCandidates(ctx, "task", wsUpdateMessageRequest{ + SessionID: "session", EntryID: "entry", + }, previous, releaser) + + require.NoError(t, err) + require.NoError(t, cleanupQueue.referenceErr) + require.NoError(t, releaser.err) + require.Equal(t, 1, releaser.calls) +} + +func TestQueueAttachmentFailureCleanupDetachesReadContext(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + cleanupQueue := &cleanupContextQueueService{QueueService: queue} + handlers.queueService = cleanupQueue + releaser := &contextRecordingReleaser{} + previous := &messagequeue.QueuedMessage{ID: "entry", TaskID: "task"} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + require.NoError(t, handlers.releaseQueuedAttachmentUpdateFailure(ctx, previous, "session", []messagequeue.MessageAttachment{{ + AttachmentID: "claimed", + }}, releaser)) + + require.NoError(t, cleanupQueue.referenceErr) + require.NoError(t, releaser.err) +} + +type failingQueueAttachmentReleaser struct{} + +func (failingQueueAttachmentReleaser) ClaimMessageAttachments(context.Context, string, string, []v1.MessageAttachment) error { + return nil +} + +func (failingQueueAttachmentReleaser) ReleaseMessageAttachments(context.Context, string, string, []v1.MessageAttachment) error { + return errors.New("release unavailable") +} + +func TestPendingAttachmentCleanupAdoptsExplicitLifecycleContext(t *testing.T) { + handlers, queue := setupQueueHandlersWithoutStart(t, nil) + current, err := queue.QueueMessage( + context.Background(), "session", "task", "current", "", messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{AttachmentID: "current"}}, + ) + require.NoError(t, err) + previous := *current + previous.Attachments = []messagequeue.MessageAttachment{{AttachmentID: "previous"}} + lifecycleCtx, cancelLifecycle := context.WithCancel(context.Background()) + + pending, err := handlers.preparePendingAttachmentCleanup( + context.Background(), + wsUpdateMessageRequest{ + SessionID: "session", EntryID: current.ID, LeaseID: "lease", OperationID: "operation", + }, + previous.TaskID, + previous.Attachments, + failingQueueAttachmentReleaser{}, + ) + require.NoError(t, err) + handlers.queuePendingAttachmentCleanup(pending) + handlers.Start(lifecycleCtx) + cancelLifecycle() + + select { + case <-handlers.attachmentCleanupCtx.Done(): + case <-time.After(100 * time.Millisecond): + t.Fatal("pending cleanup retained its lazy background context instead of the explicit lifecycle context") + } +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_race_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_race_test.go new file mode 100644 index 0000000000..0ed5a732f0 --- /dev/null +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_race_test.go @@ -0,0 +1,123 @@ +package handlers + +import ( + "context" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + v1 "github.com/kandev/kandev/pkg/api/v1" + ws "github.com/kandev/kandev/pkg/websocket" + "github.com/stretchr/testify/require" + "sync" + "testing" + "time" +) + +func TestWsUpdateMessageDoesNotReleaseAttachmentReclaimedByConcurrentEdit(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + handlers.eventBus = nil + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session-race", "task-race", "original", "", messagequeue.QueuedByUser, false, []messagequeue.MessageAttachment{{ + Type: "resource", + AttachmentID: "attachment-a", + Name: "a.txt", + MimeType: "text/plain", + }}) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + claimer := newBlockingAttachmentClaimer() + handlers.SetAttachmentClaimer(claimer) + + firstDone := make(chan updateResult, 1) + go func() { + response, callErr := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-a"), createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + "lease_id": lease.LeaseID, + "operation_id": "operation-a", + "expected_target_revision": lease.TargetRevision, + "content": "first edit", + "attachments": []messagequeue.MessageAttachment{{Type: "resource", AttachmentID: "attachment-b", Name: "b.txt", MimeType: "text/plain"}}, + })) + firstDone <- updateResult{response: response, err: callErr} + }() + <-claimer.firstReleaseStarted + secondDone := make(chan updateResult, 1) + go func() { + response, callErr := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-a"), createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + "lease_id": lease.LeaseID, + "operation_id": "operation-b", + "expected_target_revision": int64(1), + "content": "second edit", + "attachments": []messagequeue.MessageAttachment{{Type: "resource", AttachmentID: "attachment-a", Name: "a.txt", MimeType: "text/plain"}}, + })) + secondDone <- updateResult{response: response, err: callErr} + }() + select { + case <-secondDone: + t.Fatal("concurrent edit bypassed the first edit's attachment cleanup") + case <-time.After(100 * time.Millisecond): + } + close(claimer.releaseFirst) + first := <-firstDone + second := <-secondDone + require.NoError(t, first.err) + require.NoError(t, second.err) + require.Equal(t, ws.MessageTypeResponse, first.response.Type) + require.Equal(t, ws.MessageTypeResponse, second.response.Type) + require.True(t, claimer.isClaimed("attachment-a"), "attachment reclaimed by the second edit must remain claimed") +} + +type updateResult struct { + response *ws.Message + err error +} + +type blockingAttachmentClaimer struct { + mu sync.Mutex + claimed map[string]bool + releaseFirst chan struct{} + firstReleaseStarted chan struct{} + releaseCount int +} + +func newBlockingAttachmentClaimer() *blockingAttachmentClaimer { + return &blockingAttachmentClaimer{ + claimed: make(map[string]bool), + releaseFirst: make(chan struct{}), + firstReleaseStarted: make(chan struct{}), + } +} + +func (c *blockingAttachmentClaimer) ClaimMessageAttachments(_ context.Context, _ string, _ string, attachments []v1.MessageAttachment) error { + c.mu.Lock() + defer c.mu.Unlock() + for _, attachment := range attachments { + c.claimed[attachment.AttachmentID] = true + } + return nil +} + +func (c *blockingAttachmentClaimer) ReleaseMessageAttachments(_ context.Context, _ string, _ string, attachments []v1.MessageAttachment) error { + c.mu.Lock() + c.releaseCount++ + first := c.releaseCount == 1 + c.mu.Unlock() + if first { + close(c.firstReleaseStarted) + <-c.releaseFirst + } + c.mu.Lock() + defer c.mu.Unlock() + for _, attachment := range attachments { + delete(c.claimed, attachment.AttachmentID) + } + return nil +} + +func (c *blockingAttachmentClaimer) isClaimed(id string) bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.claimed[id] +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_transfer_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_transfer_test.go new file mode 100644 index 0000000000..b9edd1ad90 --- /dev/null +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_attachment_transfer_test.go @@ -0,0 +1,128 @@ +package handlers + +import ( + "context" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/kandev/kandev/internal/auth/authn" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/stretchr/testify/require" +) + +func TestLiveMissingEntryCleanupRefreshesTransferredClaimSession(t *testing.T) { + handlers, queue, db := newPersistentCleanupQueue(t, filepath.Join(t.TempDir(), "queue.db")) + defer func() { _ = db.Close() }() + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + const ( + sourceSessionID = "session-live-missing-old" + destinationSessionID = "session-live-missing-new" + taskID = "task-live-missing" + entryID = "entry-live-missing" + operationID = "cleanup-live-missing" + ) + attachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "attachment-live-missing", + Name: "live-missing.txt", MimeType: "text/plain", SizeBytes: 1, + } + require.NoError(t, queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: sourceSessionID, EntryID: entryID, OperationID: operationID, + TaskID: taskID, OwnerID: "owner", RemoveEntry: true, ClaimPending: true, + Attachments: []messagequeue.MessageAttachment{attachment}, + })) + claimer := &controlledCleanupClaimer{} + claimer.failures.Store(1) + pending := &pendingQueueAttachmentCleanup{ + key: pendingQueueAttachmentCleanupKey{ + sessionID: sourceSessionID, entryID: entryID, operationID: operationID, + }, + req: wsUpdateMessageRequest{ + SessionID: sourceSessionID, EntryID: entryID, OperationID: operationID, + }, + previous: &messagequeue.QueuedMessage{ + ID: entryID, SessionID: sourceSessionID, TaskID: taskID, + Attachments: []messagequeue.MessageAttachment{attachment}, + }, + releaser: claimer, removeEntry: true, claimPending: true, + currentSessionID: sourceSessionID, authCtx: ctx, + } + + require.Error(t, handlers.runPendingAttachmentCleanup(pending)) + require.NoError(t, queue.TransferSession(ctx, sourceSessionID, destinationSessionID)) + require.NoError(t, handlers.runPendingAttachmentCleanup(pending)) + require.Equal(t, destinationSessionID, claimer.lastSession.Load()) +} + +type transferDuringCleanupUpsertQueue struct { + *messagequeue.Service + sourceSessionID string + destinationSessionID string + transferDone chan error + transferredBeforeSave atomic.Bool +} + +func (q *transferDuringCleanupUpsertQueue) UpsertAttachmentCleanup( + ctx context.Context, + cleanup messagequeue.AttachmentCleanup, +) error { + go func() { + q.transferDone <- q.TransferSession( + context.Background(), q.sourceSessionID, q.destinationSessionID, + ) + }() + select { + case err := <-q.transferDone: + q.transferredBeforeSave.Store(true) + if err != nil { + return err + } + case <-time.After(50 * time.Millisecond): + } + return q.Service.UpsertAttachmentCleanup(ctx, cleanup) +} + +func TestCleanupPreparationHoldsAdmissionAgainstTransfer(t *testing.T) { + handlers, queue, db := newPersistentCleanupQueue(t, filepath.Join(t.TempDir(), "queue.db")) + defer func() { _ = db.Close() }() + ctx := authn.WithIdentity(context.Background(), authn.Identity{UserID: "owner"}) + entry, err := queue.QueueMessage( + ctx, "session-prepare-transfer-old", "task-prepare-transfer", "queued", "", + messagequeue.QueuedByUser, false, + []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-prepare-transfer", + Name: "prepare-transfer.txt", MimeType: "text/plain", SizeBytes: 1, + }}, + ) + require.NoError(t, err) + const destinationSessionID = "session-prepare-transfer-new" + wrapped := &transferDuringCleanupUpsertQueue{ + Service: queue, sourceSessionID: entry.SessionID, destinationSessionID: destinationSessionID, + transferDone: make(chan error, 1), + } + handlers.queueService = wrapped + _, err = handlers.preparePendingAttachmentCleanup( + ctx, + wsUpdateMessageRequest{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "prepare-transfer", + }, + entry.TaskID, + entry.Attachments, + &controlledCleanupClaimer{}, + ) + require.NoError(t, err) + require.False(t, wrapped.transferredBeforeSave.Load()) + select { + case err := <-wrapped.transferDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("session transfer remained blocked after cleanup persistence") + } + cleanup, err := queue.GetAttachmentCleanup( + ctx, entry.SessionID, entry.ID, "prepare-transfer", + ) + require.NoError(t, err) + require.NotNil(t, cleanup) + require.Equal(t, destinationSessionID, cleanup.CurrentSessionID) +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_auto_merge_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_auto_merge_test.go index 4f7383f683..788283b71e 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_handlers_auto_merge_test.go +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_auto_merge_test.go @@ -147,6 +147,15 @@ func (c *blockingQueueAttachmentClaimer) ClaimMessageAttachments(context.Context return nil } +func (*blockingQueueAttachmentClaimer) ReleaseMessageAttachments( + context.Context, + string, + string, + []v1.MessageAttachment, +) error { + return nil +} + func autoMergeQueueRequest(t *testing.T, content, attachmentID string) *ws.Message { t.Helper() return createTestMessage(t, ws.ActionMessageQueueAdd, map[string]interface{}{ diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_cancel_attachment_race_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_cancel_attachment_race_test.go new file mode 100644 index 0000000000..356c02ab30 --- /dev/null +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_cancel_attachment_race_test.go @@ -0,0 +1,154 @@ +package handlers + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + v1 "github.com/kandev/kandev/pkg/api/v1" + ws "github.com/kandev/kandev/pkg/websocket" + "github.com/stretchr/testify/require" +) + +type cancelAttachmentRaceClaimer struct { + mu sync.Mutex + claimed map[string]bool + releaseStarted chan struct{} + releaseDone chan struct{} + claimDone chan struct{} + onceRelease sync.Once + onceClaim sync.Once +} + +func newCancelAttachmentRaceClaimer() *cancelAttachmentRaceClaimer { + return &cancelAttachmentRaceClaimer{ + claimed: make(map[string]bool), + releaseStarted: make(chan struct{}), + releaseDone: make(chan struct{}), + claimDone: make(chan struct{}), + } +} + +func (c *cancelAttachmentRaceClaimer) ClaimMessageAttachments(_ context.Context, _ string, _ string, attachments []v1.MessageAttachment) error { + c.mu.Lock() + for _, attachment := range attachments { + c.claimed[attachment.AttachmentID] = true + } + c.mu.Unlock() + c.onceClaim.Do(func() { close(c.claimDone) }) + return nil +} + +func (c *cancelAttachmentRaceClaimer) ReleaseMessageAttachments(_ context.Context, _ string, _ string, attachments []v1.MessageAttachment) error { + c.onceRelease.Do(func() { + close(c.releaseStarted) + <-c.releaseDone + }) + c.mu.Lock() + defer c.mu.Unlock() + for _, attachment := range attachments { + delete(c.claimed, attachment.AttachmentID) + } + return nil +} + +func (c *cancelAttachmentRaceClaimer) isClaimed(id string) bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.claimed[id] +} + +func TestWsCancelAllDoesNotReleaseAttachmentReclaimedByConcurrentQueue(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := newCancelAttachmentRaceClaimer() + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + + entry, err := queue.QueueMessage(ctx, "session-cancel-race", "task-cancel-race", "old", "", messagequeue.QueuedByUser, false, []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-reused", Name: "old.txt", MimeType: "text/plain", + }}) + require.NoError(t, err) + claimer.claimed[entry.Attachments[0].AttachmentID] = true + + cancelDone := make(chan struct{}) + go func() { + _, _ = handlers.wsCancelAll(ctx, createTestMessage(t, ws.ActionMessageQueueCancel, map[string]string{ + "session_id": entry.SessionID, + })) + close(cancelDone) + }() + <-claimer.releaseStarted + + queueDone := make(chan struct{}) + go func() { + _, _ = handlers.wsQueueMessage(ctx, createTestMessage(t, ws.ActionMessageQueueAdd, map[string]interface{}{ + "session_id": entry.SessionID, + "task_id": entry.TaskID, + "content": "new", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-reused", Name: "new.txt", MimeType: "text/plain", + }}, + })) + close(queueDone) + }() + + select { + case <-claimer.claimDone: + case <-time.After(100 * time.Millisecond): + } + close(claimer.releaseDone) + <-cancelDone + <-queueDone + + require.True(t, claimer.isClaimed("attachment-reused"), "cancellation cleanup released an attachment used by a new queue row") + status := queue.GetStatus(ctx, entry.SessionID) + require.Len(t, status.Entries, 1) + require.Equal(t, "attachment-reused", status.Entries[0].Attachments[0].AttachmentID) +} + +func TestWsRemoveEntryDoesNotReleaseAttachmentReclaimedByConcurrentQueue(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := newCancelAttachmentRaceClaimer() + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + + entry, err := queue.QueueMessage(ctx, "session-remove-race", "task-remove-race", "old", "", messagequeue.QueuedByUser, false, []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-reused-remove", Name: "old.txt", MimeType: "text/plain", + }}) + require.NoError(t, err) + claimer.claimed[entry.Attachments[0].AttachmentID] = true + + removeDone := make(chan struct{}) + go func() { + _, _ = handlers.wsRemoveEntry(ctx, createTestMessage(t, ws.ActionMessageQueueRemove, map[string]string{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + })) + close(removeDone) + }() + <-claimer.releaseStarted + + queueDone := make(chan struct{}) + go func() { + _, _ = handlers.wsQueueMessage(ctx, createTestMessage(t, ws.ActionMessageQueueAdd, map[string]interface{}{ + "session_id": entry.SessionID, + "task_id": entry.TaskID, + "content": "new", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "attachment-reused-remove", Name: "new.txt", MimeType: "text/plain", + }}, + })) + close(queueDone) + }() + select { + case <-claimer.claimDone: + case <-time.After(100 * time.Millisecond): + } + close(claimer.releaseDone) + <-removeDone + <-queueDone + + require.True(t, claimer.isClaimed("attachment-reused-remove"), "remove cleanup released an attachment used by a new queue row") +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_cancel_count_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_cancel_count_test.go new file mode 100644 index 0000000000..58dd4eee91 --- /dev/null +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_cancel_count_test.go @@ -0,0 +1,38 @@ +package handlers + +import ( + "context" + "encoding/json" + "testing" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + ws "github.com/kandev/kandev/pkg/websocket" + "github.com/stretchr/testify/require" +) + +// queueServiceWithoutBatch preserves the legacy handler surface so the +// fallback cancellation path is exercised without CancelAllWithEntries. +type queueServiceWithoutBatch struct { + QueueService +} + +func TestWsCancelAllFallbackReportsActualRemovedCount(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + handlers.queueService = queueServiceWithoutBatch{QueueService: queue} + ctx := context.Background() + + _, err := queue.QueueMessage(ctx, "session-cancel-fallback", "task", "queued", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + + response, err := handlers.wsCancelAll(ctx, createTestMessage(t, ws.ActionMessageQueueCancel, map[string]string{ + "session_id": "session-cancel-fallback", + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + + var payload struct { + Removed int `json:"removed"` + } + require.NoError(t, json.Unmarshal(response.Payload, &payload)) + require.Equal(t, 1, payload.Removed) +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_edit_lease_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_edit_lease_test.go new file mode 100644 index 0000000000..75c5b65c58 --- /dev/null +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_edit_lease_test.go @@ -0,0 +1,105 @@ +package handlers + +import ( + "context" + "testing" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + ws "github.com/kandev/kandev/pkg/websocket" + "github.com/stretchr/testify/require" +) + +func TestWsUpdateMessageDoesNotReleaseUnclaimedAttachmentsOnLeaseConflict(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + claimer := &recordingQueueAttachmentClaimer{} + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session", "task", "original", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + firstLease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + require.NoError(t, queue.EndEdit(ctx, entry.SessionID, entry.ID, firstLease.LeaseID, "connection-a")) + secondLease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-b") + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ws.WithConnectionID(ctx, "connection-a"), createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + "lease_id": firstLease.LeaseID, + "operation_id": "operation-1", + "expected_target_revision": firstLease.TargetRevision, + "content": "edited", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "staged", Name: "staged.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + require.Empty(t, claimer.claims) + require.Empty(t, claimer.releases) + require.NoError(t, queue.EndEdit(ctx, entry.SessionID, entry.ID, secondLease.LeaseID, "connection-b")) +} + +func TestWsBeginEditReturnsEntryNotFoundForDrainedEntry(t *testing.T) { + handlers, queue := setupQueueHandlers(t) + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session", "task", "original", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + _, ok := queue.TakeQueued(ctx, entry.SessionID) + require.True(t, ok) + + response, err := handlers.wsBeginEdit( + ws.WithConnectionID(ctx, "connection"), + createTestMessage(t, ws.ActionMessageQueueEditBegin, map[string]interface{}{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + }), + ) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeError, response.Type) + require.Equal(t, queueErrorCodeEntryNotFound, parseError(t, response).Code) +} + +func TestWsEndEditDispatchesOnlyAfterSuccessfulSave(t *testing.T) { + drainer := &mockQueueDrainer{} + handlers, queue := setupQueueHandlersWithDrainer(t, drainer) + ctx := context.Background() + entry, err := queue.QueueMessage(ctx, "session", "task", "original", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := queue.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + _, err = queue.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection", lease.TargetRevision, "edited", nil, nil) + require.NoError(t, err) + + response, err := handlers.wsEndEdit( + ws.WithConnectionID(ctx, "connection"), + createTestMessage(t, ws.ActionMessageQueueEditEnd, map[string]interface{}{ + "session_id": entry.SessionID, + "entry_id": entry.ID, + "lease_id": lease.LeaseID, + "dispatch_if_auto_run": true, + }), + ) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Equal(t, 1, drainer.autoRunCalls) + require.Equal(t, entry.SessionID, drainer.sessionID) + + second, err := queue.QueueMessage(ctx, entry.SessionID, entry.TaskID, "not saved", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + secondLease, err := queue.BeginEdit(ctx, second.SessionID, second.ID, "connection") + require.NoError(t, err) + response, err = handlers.wsEndEdit( + ws.WithConnectionID(ctx, "connection"), + createTestMessage(t, ws.ActionMessageQueueEditEnd, map[string]interface{}{ + "session_id": second.SessionID, + "entry_id": second.ID, + "lease_id": secondLease.LeaseID, + "dispatch_if_auto_run": true, + }), + ) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, response.Type) + require.Equal(t, 1, drainer.autoRunCalls) +} diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_publish_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_publish_test.go index 1fd06b41d1..173212c28f 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_handlers_publish_test.go +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_publish_test.go @@ -16,10 +16,14 @@ import ( // capturingQueueEventBus records the last message.queue.status_changed payload // so tests can assert the task_id enrichment on the published event. type capturingQueueEventBus struct { - lastData map[string]interface{} + lastData map[string]interface{} + lastSubject string + lastCtx context.Context } -func (m *capturingQueueEventBus) Publish(_ context.Context, _ string, event *bus.Event) error { +func (m *capturingQueueEventBus) Publish(ctx context.Context, subject string, event *bus.Event) error { + m.lastCtx = ctx + m.lastSubject = subject if data, ok := event.Data.(map[string]interface{}); ok { m.lastData = data } @@ -51,9 +55,61 @@ func setupQueueHandlersWithResolver( events := &capturingQueueEventBus{} svc := messagequeue.NewServiceMemory(log) handlers := NewQueueHandlers(svc, events, log, nil, allowQueueAccess{}, resolver) + handlers.Start(context.Background()) + t.Cleanup(handlers.Stop) return handlers, svc, events } +type blockingQueueEventBus struct { + capturingQueueEventBus + entered chan struct{} + release chan struct{} +} + +func (b *blockingQueueEventBus) Publish(context.Context, string, *bus.Event) error { + b.entered <- struct{}{} + <-b.release + return nil +} + +func TestPublishStatusSerializesSnapshotAndPublish(t *testing.T) { + log, err := logger.NewLogger(logger.LoggingConfig{ + Level: "error", + Format: "console", + OutputPath: "stderr", + }) + require.NoError(t, err) + events := &blockingQueueEventBus{ + entered: make(chan struct{}, 2), + release: make(chan struct{}), + } + svc := messagequeue.NewServiceMemory(log) + handlers := NewQueueHandlers(svc, events, log, nil, allowQueueAccess{}, nil) + ctx := context.Background() + + firstDone := make(chan struct{}) + go func() { + handlers.publishStatus(ctx, "s1") + close(firstDone) + }() + <-events.entered + + secondDone := make(chan struct{}) + go func() { + handlers.publishStatus(ctx, "s1") + close(secondDone) + }() + select { + case <-events.entered: + t.Fatal("second queue status reached the event bus before the first publish completed") + case <-time.After(100 * time.Millisecond): + } + + close(events.release) + <-firstDone + <-secondDone +} + func TestPublishStatusIncludesTaskIDWhenResolvable(t *testing.T) { handlers, svc, events := setupQueueHandlersWithResolver(t, func(_ context.Context, sessionID string) (string, error) { if sessionID == "s1" { @@ -77,6 +133,17 @@ func TestPublishStatusIncludesTaskIDWhenResolvable(t *testing.T) { require.Equal(t, true, events.lastData["merge_enabled"]) } +func TestPublishStatusDetachesCancelledRequest(t *testing.T) { + handlers, _, events := setupQueueHandlersWithResolver(t, nil) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + handlers.publishStatus(ctx, "s1") + + require.NotNil(t, events.lastCtx) + require.NoError(t, events.lastCtx.Err(), "queue status publication must not inherit a cancelled request") +} + func TestWsQueueMessagePublishesUserPromptActivity(t *testing.T) { handlers, _, events := setupQueueHandlersWithResolver(t, func(_ context.Context, sessionID string) (string, error) { if sessionID == "s1" { diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_reorder_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_reorder_test.go index e5f6a898fd..e72cc7cf02 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_handlers_reorder_test.go +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_reorder_test.go @@ -106,6 +106,25 @@ func TestWsReorder(t *testing.T) { require.Len(t, status.Entries, 1) assert.Equal(t, second.ID, status.Entries[0].ID) }) + t.Run("rejects reorder while an entry is leased for editing", func(t *testing.T) { + handlers, svc := setupQueueHandlers(t) + ctx := context.Background() + first, err := svc.QueueMessage(ctx, "s", "t", "first", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + second, err := svc.QueueMessage(ctx, "s", "t", "second", "", messagequeue.QueuedByUser, false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, "s", first.ID, "connection-a") + require.NoError(t, err) + + response, err := handlers.wsReorder(ctx, createTestMessage(t, ws.ActionMessageQueueReorder, map[string]interface{}{ + "session_id": "s", + "ordered_ids": []string{second.ID, first.ID}, + })) + require.NoError(t, err) + assert.Equal(t, ws.MessageTypeError, response.Type) + assert.Equal(t, "edit_conflict", parseError(t, response).Code) + assert.Equal(t, first.ID, svc.GetStatus(ctx, "s").Entries[0].ID) + }) t.Run("denies when the session is not authorized", func(t *testing.T) { log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console", OutputPath: "stderr"}) diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_send_now_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_send_now_test.go index cd64ef6a2e..84b1e2476e 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_handlers_send_now_test.go +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_send_now_test.go @@ -91,6 +91,7 @@ func TestWsSendNowMapsStableErrors(t *testing.T) { {name: "empty", err: orchestrator.ErrSendNowQueueEmpty, code: queueErrorCodeSendNowQueueEmpty}, {name: "changed", err: orchestrator.ErrSendNowQueueChanged, code: queueErrorCodeSendNowQueueChanged}, {name: "conflict", err: orchestrator.ErrSendNowConflict, code: queueErrorCodeSendNowConflict}, + {name: "edit conflict", err: orchestrator.ErrSendNowEditConflict, code: "edit_conflict"}, {name: "turn", err: orchestrator.ErrSendNowTurnChanged, code: queueErrorCodeSendNowTurnChanged}, {name: "attachments", err: messagequeue.ErrSendNowAttachmentOverflow, code: queueErrorCodeSendNowAttachmentOverflow}, {name: "references", err: messagequeue.ErrSendNowReferenceOverflow, code: queueErrorCodeSendNowReferenceOverflow}, diff --git a/apps/backend/internal/orchestrator/handlers/queue_handlers_test.go b/apps/backend/internal/orchestrator/handlers/queue_handlers_test.go index 70f89f40a2..2384ea90a0 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_handlers_test.go +++ b/apps/backend/internal/orchestrator/handlers/queue_handlers_test.go @@ -5,9 +5,6 @@ import ( "encoding/base64" "encoding/json" "errors" - "testing" - "time" - "github.com/kandev/kandev/internal/common/logger" "github.com/kandev/kandev/internal/entityrefs" "github.com/kandev/kandev/internal/events/bus" @@ -18,6 +15,8 @@ import ( ws "github.com/kandev/kandev/pkg/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "testing" + "time" ) // mockEventBus is a no-op event bus for handler tests. @@ -125,6 +124,11 @@ func (m *mockQueueDrainer) DrainQueuedMessage(_ context.Context, sessionID strin m.sessionID = sessionID return m.drained, m.err } +func (m *mockQueueDrainer) DrainQueuedMessageIfAutoRun(_ context.Context, sessionID string) (bool, error) { + m.autoRunCalls++ + m.sessionID = sessionID + return m.autoRunResult, m.autoRunErr +} func (m *mockQueueDrainer) SetQueueAutoRun(_ context.Context, sessionID string, enabled bool) (bool, bool, error) { m.autoRunCalls++ @@ -138,6 +142,13 @@ func setupQueueHandlers(t *testing.T) (*QueueHandlers, *messagequeue.Service) { } func setupQueueHandlersWithDrainer(t *testing.T, drainer QueueDrainer) (*QueueHandlers, *messagequeue.Service) { + t.Helper() + handlers, svc := setupQueueHandlersWithoutStart(t, drainer) + handlers.Start(context.Background()) + return handlers, svc +} + +func setupQueueHandlersWithoutStart(t *testing.T, drainer QueueDrainer) (*QueueHandlers, *messagequeue.Service) { t.Helper() log, err := logger.NewLogger(logger.LoggingConfig{ Level: "error", @@ -147,7 +158,9 @@ func setupQueueHandlersWithDrainer(t *testing.T, drainer QueueDrainer) (*QueueHa require.NoError(t, err) svc := messagequeue.NewServiceMemory(log) svc.SetAutoMergeEnabled(false) - return NewQueueHandlers(svc, &mockEventBus{}, log, drainer, allowQueueAccess{}, nil), svc + handlers := NewQueueHandlers(svc, &mockEventBus{}, log, drainer, allowQueueAccess{}, nil) + t.Cleanup(handlers.Stop) + return handlers, svc } func setupQueueHandlersWithValidator(t *testing.T, validator entityrefs.SubmissionValidator) (*QueueHandlers, *messagequeue.Service) { @@ -160,7 +173,10 @@ func setupQueueHandlersWithValidator(t *testing.T, validator entityrefs.Submissi require.NoError(t, err) svc := messagequeue.NewServiceMemory(log) svc.SetAutoMergeEnabled(false) - return NewQueueHandlers(svc, &mockEventBus{}, log, nil, allowQueueAccess{}, nil, validator), svc + handlers := NewQueueHandlers(svc, &mockEventBus{}, log, nil, allowQueueAccess{}, nil, validator) + handlers.Start(context.Background()) + t.Cleanup(handlers.Stop) + return handlers, svc } func createTestMessage(t *testing.T, action string, payload interface{}) *ws.Message { @@ -742,6 +758,28 @@ func TestWsUpdateMessage(t *testing.T) { assert.Equal(t, []string{"new-attachment"}, claimer.claims) assert.Equal(t, []string{"old-attachment"}, claimer.releases) }) + t.Run("preserves an attachment still referenced by another queue entry", func(t *testing.T) { + handlers, svc := setupQueueHandlers(t) + claimer := &recordingQueueAttachmentClaimer{} + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + attachment := messagequeue.MessageAttachment{ + Type: "resource", AttachmentID: "shared-attachment", Name: "shared.txt", MimeType: "text/plain", + } + first, err := svc.QueueMessage(ctx, "s", "task-1", "first", "", "u", false, []messagequeue.MessageAttachment{attachment}) + require.NoError(t, err) + _, err = svc.QueueMessage(ctx, "s", "task-1", "second", "", "u", false, []messagequeue.MessageAttachment{attachment}) + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ctx, + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": "s", "entry_id": first.ID, "content": "edited", "user_id": "u", + "attachments": []messagequeue.MessageAttachment{}, + })) + require.NoError(t, err) + assert.Equal(t, ws.MessageTypeResponse, response.Type) + assert.Empty(t, claimer.releases) + }) t.Run("releases only newly claimed attachments when replacement fails", func(t *testing.T) { handlers, svc := setupQueueHandlers(t) @@ -767,6 +805,27 @@ func TestWsUpdateMessage(t *testing.T) { assert.Equal(t, []string{"new-attachment"}, claimer.releases) }) + t.Run("releases attachments when claim reports an error", func(t *testing.T) { + handlers, svc := setupQueueHandlers(t) + claimer := &recordingQueueAttachmentClaimer{claimErr: errors.New("claim failed")} + handlers.SetAttachmentClaimer(claimer) + ctx := context.Background() + queued, err := svc.QueueMessage(ctx, "s", "task-1", "original", "", "u", false, nil) + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ctx, + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": "s", "entry_id": queued.ID, "content": "edited", "user_id": "u", + "attachments": []messagequeue.MessageAttachment{{ + Type: "resource", AttachmentID: "new-attachment", Name: "new.txt", MimeType: "text/plain", + }}, + })) + require.NoError(t, err) + assert.Equal(t, ws.MessageTypeError, response.Type) + assert.Equal(t, []string{"new-attachment"}, claimer.claims) + assert.Equal(t, []string{"new-attachment"}, claimer.releases) + }) + t.Run("updates an entry", func(t *testing.T) { handlers, svc := setupQueueHandlers(t) ctx := context.Background() @@ -1017,6 +1076,30 @@ func TestWsUpdateMessage(t *testing.T) { require.Len(t, entries, 1) assert.Equal(t, "lifecycle prompt", entries[0].Content) }) + + t.Run("rejects user_id impersonating the move-task identity", func(t *testing.T) { + handlers, svc := setupQueueHandlers(t) + ctx := context.Background() + queued, err := svc.QueueMessageWithMetadata( + ctx, "s", "t", "handoff", "", messagequeue.QueuedByMoveTask, false, nil, + map[string]interface{}{messagequeue.MetadataDeferredMoveID: "move-1"}, + ) + require.NoError(t, err) + + response, err := handlers.wsUpdateMessage(ctx, + createTestMessage(t, ws.ActionMessageQueueUpdate, map[string]interface{}{ + "session_id": "s", + "entry_id": queued.ID, + "content": "forged handoff", + "user_id": messagequeue.QueuedByMoveTask, + })) + require.NoError(t, err) + assert.Equal(t, ws.MessageTypeError, response.Type) + + entries := svc.GetStatus(ctx, "s").Entries + require.Len(t, entries, 1) + assert.Equal(t, "handoff", entries[0].Content) + }) } func TestWsRemoveEntry(t *testing.T) { diff --git a/apps/backend/internal/orchestrator/handlers/queue_operations_handlers.go b/apps/backend/internal/orchestrator/handlers/queue_operations_handlers.go index a98b8c9524..4c046284b0 100644 --- a/apps/backend/internal/orchestrator/handlers/queue_operations_handlers.go +++ b/apps/backend/internal/orchestrator/handlers/queue_operations_handlers.go @@ -38,6 +38,17 @@ func (h *QueueHandlers) wsCancelAll(ctx context.Context, msg *ws.Message) (*ws.M identity := messagequeue.QueueSessionIdentity{ TaskID: req.TaskID, SessionID: req.SessionID, SessionIncarnationID: req.SessionIncarnationID, } + if !h.requiresQueueIdentity() { + removed, err := h.cancelAllLegacy(ctx, req.SessionID) + if err != nil { + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, err.Error(), nil) + } + h.publishStatus(ctx, req.SessionID) + return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{ + fieldSessionID: req.SessionID, + "removed": removed, + }) + } var removal *messagequeue.QueueRemovalResult var removed int var err error @@ -302,6 +313,8 @@ func (h *QueueHandlers) sendNowErrorResponse(msg *ws.Message, sessionID string, return ws.NewError(msg.ID, msg.Action, queueErrorCodeSendNowQueueEmpty, "Queue is empty", nil) case errors.Is(err, orchestrator.ErrSendNowQueueChanged): return ws.NewError(msg.ID, msg.Action, queueErrorCodeSendNowQueueChanged, "Queue changed before Send Now could start", nil) + case errors.Is(err, orchestrator.ErrSendNowEditConflict): + return ws.NewError(msg.ID, msg.Action, "edit_conflict", "Queue entry is being edited by another view", nil) case errors.Is(err, orchestrator.ErrSendNowConflict): return ws.NewError(msg.ID, msg.Action, queueErrorCodeSendNowConflict, "Another cancellation or Send Now operation is in progress", nil) case errors.Is(err, orchestrator.ErrSendNowTurnChanged): @@ -385,6 +398,16 @@ func (h *QueueHandlers) wsRemoveEntry(ctx context.Context, msg *ws.Message) (*ws if req.EntryID == "" { return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "entry_id is required", nil) } + if !h.requiresQueueIdentity() { + if err := h.removeEntryLegacy(ctx, req.SessionID, req.EntryID); err != nil { + if errors.Is(err, messagequeue.ErrEntryNotFound) { + return ws.NewError(msg.ID, msg.Action, queueErrorCodeEntryNotFound, "Queue entry is no longer pending", nil) + } + return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, err.Error(), nil) + } + h.publishStatus(ctx, req.SessionID) + return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{fieldEntryID: req.EntryID}) + } identity := messagequeue.QueueSessionIdentity{ TaskID: req.TaskID, SessionID: req.SessionID, SessionIncarnationID: req.SessionIncarnationID, @@ -418,6 +441,110 @@ func (h *QueueHandlers) wsRemoveEntry(ctx context.Context, msg *ws.Message) (*ws return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{fieldEntryID: req.EntryID}) } +// cancelAllLegacy preserves the attachment cleanup contract for the +// pre-incarnation queue API. Identity-bound callers use the atomic removal +// result returned by CancelAllForSession above. +// +//nolint:cyclop,gocognit // Legacy cancellation preserves attachment cleanup across two queue service contracts. +func (h *QueueHandlers) cancelAllLegacy(ctx context.Context, sessionID string) (int, error) { + var removed int + var removedEntries []messagequeue.QueuedMessage + var countRemovedEntries bool + cancel := func(cancelCtx context.Context) error { + prepared := make(map[string]*pendingQueueAttachmentCleanup) + status := h.queueService.GetStatus(cancelCtx, sessionID) + if status != nil { + for i := range status.Entries { + entry := &status.Entries[i] + if entry.IsReservedInFlight() { + continue + } + pending, prepareErr := h.prepareEntryRemovalCleanup(cancelCtx, entry, "cancel") + if prepareErr != nil { + for _, earlier := range prepared { + h.settlePendingAttachmentCleanup(earlier, nil) + } + return prepareErr + } + if pending != nil { + prepared[entry.ID] = pending + } + } + } + + var err error + if batchCanceller, ok := h.queueService.(queueBatchCanceller); ok { + removedEntries, err = batchCanceller.CancelAllWithEntries(cancelCtx, sessionID) + countRemovedEntries = true + } else { + removed, err = h.queueService.CancelAll(cancelCtx, sessionID) + if err == nil && status != nil { + removedEntries = status.Entries + } + } + if err != nil { + for _, pending := range prepared { + h.settlePendingAttachmentCleanup(pending, nil) + } + return err + } + for i := range removedEntries { + if removedEntries[i].IsReservedInFlight() { + continue + } + if countRemovedEntries { + removed++ + } + h.releaseQueuedAttachments(cancelCtx, &removedEntries[i], prepared[removedEntries[i].ID]) + } + return nil + } + if admission, ok := h.queueService.(queueEditAdmissionController); ok { + if err := admission.WithSessionAdmission(ctx, sessionID, cancel); err != nil { + return 0, err + } + return removed, nil + } + if err := cancel(ctx); err != nil { + return 0, err + } + return removed, nil +} + +// removeEntryLegacy preserves durable attachment cleanup for the +// pre-incarnation queue API. The cleanup obligation is prepared before the +// row is removed and settled only after the removal commits. +func (h *QueueHandlers) removeEntryLegacy(ctx context.Context, sessionID, entryID string) error { + var err error + remove := func(removeCtx context.Context) error { + entry, getErr := h.queueService.GetEntry(removeCtx, sessionID, entryID) + if getErr != nil { + return getErr + } + pending, prepareErr := h.prepareEntryRemovalCleanup(removeCtx, entry, "remove") + if prepareErr != nil { + return prepareErr + } + if remover, ok := h.queueService.(queueEntryRemover); ok { + entry, err = remover.RemoveEntryWithEntry(removeCtx, sessionID, entryID) + } else { + err = h.queueService.RemoveEntry(removeCtx, sessionID, entryID) + } + if err != nil { + h.settlePendingAttachmentCleanup(pending, nil) + return err + } + h.releaseQueuedAttachments(removeCtx, entry, pending) + return nil + } + if admission, ok := h.queueService.(queueEditAdmissionController); ok { + err = admission.WithSessionAdmission(ctx, sessionID, remove) + } else { + err = remove(ctx) + } + return err +} + func (h *QueueHandlers) releaseQueueRemovalAttachments(ctx context.Context, removal *messagequeue.QueueRemovalResult) { releaser, ok := h.attachmentClaimer.(QueueAttachmentReleaser) if !ok || removal == nil || len(removal.Removed) == 0 { @@ -589,6 +716,9 @@ func (h *QueueHandlers) wsReorder(ctx context.Context, msg *ws.Message) (*ws.Mes if errors.Is(err, messagequeue.ErrQueueChanged) { return ws.NewError(msg.ID, msg.Action, queueErrorCodeQueueChanged, "Queue changed before the reorder could be applied", nil) } + if errors.Is(err, messagequeue.ErrEditConflict) { + return h.queueEditLeaseError(msg, err), nil + } if isQueueIdentityError(err) { return queueAccessDeniedResponse(msg), nil } @@ -609,6 +739,20 @@ func (h *QueueHandlers) publishStatus(ctx context.Context, sessionID string, adm if h.eventBus == nil { return } + // A committed queue mutation still needs an authoritative snapshot when + // the initiating request has already been cancelled. + ctx = context.WithoutCancel(ctx) + if admission, ok := h.queueService.(queueEditAdmissionController); ok { + _ = admission.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + h.publishStatusSnapshot(admittedCtx, sessionID, admitted...) + return nil + }) + return + } + h.publishStatusSnapshot(ctx, sessionID, admitted...) +} + +func (h *QueueHandlers) publishStatusSnapshot(ctx context.Context, sessionID string, admitted ...*messagequeue.QueuedMessage) { status := h.queueService.GetStatus(ctx, sessionID) eventData := map[string]interface{}{ fieldSessionID: sessionID, diff --git a/apps/backend/internal/orchestrator/messagequeue/edit_lease_lifecycle_test.go b/apps/backend/internal/orchestrator/messagequeue/edit_lease_lifecycle_test.go new file mode 100644 index 0000000000..7fe8085004 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/edit_lease_lifecycle_test.go @@ -0,0 +1,23 @@ +package messagequeue + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestInvalidateEditLeasesForTaskRejectsStaleUpdates(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lifecycle", "task-lifecycle", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + svc.InvalidateEditLeasesForTask(entry.TaskID) + + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection-a", lease.TargetRevision, "stale edit", nil, nil) + require.ErrorIs(t, err, ErrEditLeaseNotFound) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/edit_lease_mutation_test.go b/apps/backend/internal/orchestrator/messagequeue/edit_lease_mutation_test.go new file mode 100644 index 0000000000..445ae73d69 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/edit_lease_mutation_test.go @@ -0,0 +1,324 @@ +package messagequeue + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestEditLeasePreventsAutomaticMergeIntoTarget(t *testing.T) { + svc := setupService(t) + svc.SetAutoMergeEnabled(true) + ctx := context.Background() + first, err := svc.QueueMessage(ctx, "session-lease-auto-merge", "task", "first", "", "user", false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, first.SessionID, first.ID, "connection-a") + require.NoError(t, err) + + _, err = svc.QueueMessage(ctx, first.SessionID, "task", "second", "", "user", false, nil) + require.NoError(t, err) + + entries := svc.GetStatus(ctx, first.SessionID).Entries + require.Len(t, entries, 2) + require.Equal(t, "first", entries[0].Content) + require.Equal(t, "second", entries[1].Content) +} + +func TestEditLeasePreventsManualMergeIntoTarget(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + first, err := svc.QueueMessage(ctx, "session-lease-manual-merge", "task", "first", "", "user", false, nil) + require.NoError(t, err) + second, err := svc.QueueMessage(ctx, first.SessionID, "task", "second", "", "user", false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, first.SessionID, first.ID, "connection-a") + require.NoError(t, err) + + _, err = svc.MergeIntoAbove(ctx, first.SessionID, second.ID, "user") + require.ErrorIs(t, err, ErrEditConflict) + + entries := svc.GetStatus(ctx, first.SessionID).Entries + require.Len(t, entries, 2) + require.Equal(t, "first", entries[0].Content) + require.Equal(t, "second", entries[1].Content) + require.False(t, errors.Is(err, ErrEntryNotFound)) +} +func TestEditLeaseBlocksReorderOfTarget(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + first, err := svc.QueueMessage(ctx, "session-lease-reorder", "task", "first", "", "user", false, nil) + require.NoError(t, err) + second, err := svc.QueueMessage(ctx, first.SessionID, "task", "second", "", "user", false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, first.SessionID, first.ID, "connection-a") + require.NoError(t, err) + + err = svc.ReorderEntries(ctx, first.SessionID, []string{second.ID, first.ID}) + require.ErrorIs(t, err, ErrEditConflict) + + entries := svc.GetStatus(ctx, first.SessionID).Entries + require.Len(t, entries, 2) + require.Equal(t, first.ID, entries[0].ID) + require.Equal(t, second.ID, entries[1].ID) +} + +func TestEditLeaseInvalidatedBySessionReplacement(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-replace", "task", "original", "", "user", false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + replacement := *entry + replacement.Content = "restored" + require.NoError(t, svc.RestoreSession(ctx, entry.SessionID, []QueuedMessage{replacement}, nil)) + + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection-a", lease.TargetRevision, "stale edit", nil, nil) + require.ErrorIs(t, err, ErrEditLeaseNotFound) + require.Equal(t, "restored", svc.GetStatus(ctx, entry.SessionID).Entries[0].Content) +} + +func TestEditLeaseBlocksAppendIntoTarget(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + first, err := svc.QueueMessage(ctx, "session-lease-append", "task", "first", "", "user", false, nil) + require.NoError(t, err) + second, err := svc.QueueMessage(ctx, first.SessionID, "task", "second", "", "user", false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, first.SessionID, second.ID, "connection-a") + require.NoError(t, err) + + _, _, err = svc.AppendContent(ctx, first.SessionID, "task", " appended", "", "user", false, nil) + require.ErrorIs(t, err, ErrEditConflict) + require.Equal(t, "second", svc.GetStatus(ctx, first.SessionID).Entries[1].Content) +} + +func TestEditLeaseAllowsAppendToOtherEntry(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + first, err := svc.QueueMessage(ctx, "session-lease-append-other", "task", "first", "", "user", false, nil) + require.NoError(t, err) + _, err = svc.QueueMessage(ctx, first.SessionID, "task", "second", "", "user", false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, first.SessionID, first.ID, "connection-a") + require.NoError(t, err) + + merged, appended, err := svc.AppendContent(ctx, first.SessionID, "task", " appended", "", "user", false, nil) + require.NoError(t, err) + require.True(t, appended) + require.Equal(t, "second\n\n---\n\n appended", merged.Content) +} +func TestPurgeTaskInvalidatesEditLeases(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-purge", "task-purge", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + removed, err := svc.PurgeTask(ctx, entry.TaskID) + require.NoError(t, err) + require.Equal(t, 1, removed) + require.Empty(t, svc.editLeases) +} +func TestEditRevisionStateFollowsEntryLifecycle(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + + entry, err := svc.QueueMessage(ctx, "session-edit-delete", "task-edit-delete", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection-a", lease.TargetRevision, "edited", nil, nil) + require.NoError(t, err) + require.Contains(t, svc.editRevisions, svc.editLeaseKey(entry.SessionID, entry.ID)) + require.NoError(t, svc.EndEdit(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection-a")) + require.NoError(t, svc.RemoveEntry(ctx, entry.SessionID, entry.ID)) + require.NotContains(t, svc.editRevisions, svc.editLeaseKey(entry.SessionID, entry.ID)) + require.NotContains(t, svc.editRevisionTaskIDs, svc.editLeaseKey(entry.SessionID, entry.ID)) + + entry, err = svc.QueueMessage(ctx, "session-edit-transfer", "task-edit-transfer", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-b") + require.NoError(t, err) + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-2", "connection-b", lease.TargetRevision, "edited", nil, nil) + require.NoError(t, err) + require.NoError(t, svc.TransferSession(ctx, entry.SessionID, "session-edit-transfer-new")) + require.NotContains(t, svc.editRevisions, svc.editLeaseKey(entry.SessionID, entry.ID)) + + entry, err = svc.QueueMessage(ctx, "session-edit-purge", "task-edit-purge", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-c") + require.NoError(t, err) + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-3", "connection-c", lease.TargetRevision, "edited", nil, nil) + require.NoError(t, err) + _, err = svc.PurgeTask(ctx, entry.TaskID) + require.NoError(t, err) + require.NotContains(t, svc.editRevisions, svc.editLeaseKey(entry.SessionID, entry.ID)) + require.NotContains(t, svc.editRevisionTaskIDs, svc.editLeaseKey(entry.SessionID, entry.ID)) +} +func TestPurgeTaskPreservesOtherEditRevisions(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + live, err := svc.QueueMessage(ctx, "session-live-edit", "task-live", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + other, err := svc.QueueMessage(ctx, "session-other-edit", "task-other", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, live.SessionID, live.ID, "connection-a") + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, other.SessionID, other.ID, "connection-b") + require.NoError(t, err) + + revision, err := svc.UpdateMessageWithLease(ctx, live.SessionID, live.ID, lease.LeaseID, + "operation-live", "connection-a", lease.TargetRevision, "edited", nil, nil) + require.NoError(t, err) + require.Equal(t, int64(1), revision) + + _, err = svc.PurgeTask(ctx, other.TaskID) + require.NoError(t, err) + + revision, err = svc.UpdateMessageWithLease(ctx, live.SessionID, live.ID, lease.LeaseID, + "operation-live-2", "connection-a", revision, "edited again", nil, nil) + require.NoError(t, err) + require.Equal(t, int64(2), revision) +} +func TestLeaseUpdateFinalizerRunsBeforeSessionTransfer(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-finalize-old", "task-finalize", "before", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + transferDone := make(chan error, 1) + _, err = svc.UpdateMessageWithLeaseAfterValidationAndFinalize( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation-finalize", "connection-a", + lease.TargetRevision, "after", nil, nil, nil, nil, + func(context.Context, *QueuedMessage) error { + go func() { + transferDone <- svc.TransferSession(context.Background(), entry.SessionID, "session-finalize-new") + }() + select { + case err := <-transferDone: + t.Fatalf("session transfer completed before finalizer returned: %v", err) + case <-time.After(25 * time.Millisecond): + } + return nil + }, + ) + require.NoError(t, err) + require.NoError(t, <-transferDone) + require.Empty(t, svc.GetStatus(ctx, entry.SessionID).Entries) + require.Equal(t, "after", svc.GetStatus(ctx, "session-finalize-new").Entries[0].Content) +} + +func TestLeaseUpdateFinalizerReplayKeepsOriginalPreUpdateSnapshot(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + original := MessageAttachment{AttachmentID: "original"} + replacement := MessageAttachment{AttachmentID: "replacement"} + entry, err := svc.QueueMessage( + ctx, "session-finalize-replay", "task-finalize", "before", "", QueuedByUser, false, + []MessageAttachment{original}, + ) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + finalizeFailure := errors.New("finalize unavailable") + finalizeCalls := 0 + finalize := func(_ context.Context, previous *QueuedMessage) error { + finalizeCalls++ + require.Equal(t, []MessageAttachment{original}, previous.Attachments) + if finalizeCalls == 1 { + return finalizeFailure + } + return nil + } + + revision, err := svc.UpdateMessageWithLeaseAfterValidationAndFinalize( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation-finalize-replay", "connection-a", + lease.TargetRevision, "after", []MessageAttachment{replacement}, nil, nil, nil, finalize, + ) + require.ErrorIs(t, err, finalizeFailure) + require.Equal(t, int64(1), revision) + require.Equal(t, []MessageAttachment{replacement}, svc.GetStatus(ctx, entry.SessionID).Entries[0].Attachments) + + revision, err = svc.UpdateMessageWithLeaseAfterValidationAndFinalize( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation-finalize-replay", "connection-a", + lease.TargetRevision, "after", []MessageAttachment{replacement}, nil, nil, nil, finalize, + ) + require.NoError(t, err) + require.Equal(t, int64(1), revision) + require.Equal(t, 2, finalizeCalls) +} +func TestEditLeaseChecksPhysicalHeadAndTail(t *testing.T) { + for _, factory := range autoRunRepositoryFactories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + svc := newAutoMergeTestServiceWithRepository(t, repo, 10) + svc.SetAutoMergeEnabled(false) + ctx := context.Background() + + t.Run("reserved head does not expose a leased visible row", func(t *testing.T) { + head := &QueuedMessage{ + SessionID: "session-physical-head", TaskID: "task-head", + Content: "head", QueuedBy: QueuedByWorkflow, + Metadata: map[string]interface{}{MetadataLifecycleDurable: true}, + } + require.NoError(t, repo.Insert(ctx, head, 0)) + _, err := svc.QueueMessage(ctx, head.SessionID, "task-user", "user", "", QueuedByUser, false, nil) + require.NoError(t, err) + reserved, ok := svc.ReserveQueued(ctx, head.SessionID) + require.True(t, ok) + require.Equal(t, head.ID, reserved.ID) + user := svc.GetStatus(ctx, head.SessionID).Entries[0] + _, err = svc.BeginEdit(ctx, user.SessionID, user.ID, "connection") + require.NoError(t, err) + + recovered, ok := svc.ReserveQueued(ctx, head.SessionID) + require.True(t, ok) + require.Equal(t, head.ID, recovered.ID) + }) + + t.Run("reserved tail does not expose a leased visible row", func(t *testing.T) { + user, err := svc.QueueMessage(ctx, "session-physical-tail", "task-user", "user", "", QueuedByUser, false, nil) + require.NoError(t, err) + tail := &QueuedMessage{ + SessionID: user.SessionID, TaskID: "task-tail", + Content: "tail", QueuedBy: QueuedByWorkflow, + Metadata: map[string]interface{}{MetadataLifecycleDurable: true}, + } + require.NoError(t, repo.Insert(ctx, tail, 0)) + taken, ok := svc.TakeQueued(ctx, user.SessionID) + require.True(t, ok) + require.Equal(t, user.ID, taken.ID) + reserved, ok := svc.ReserveQueued(ctx, user.SessionID) + require.True(t, ok) + require.Equal(t, tail.ID, reserved.ID) + restored, err := svc.RestoreMessage(ctx, taken) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, restored.SessionID, restored.ID, "connection") + require.NoError(t, err) + + inserted, appended, err := svc.AppendContent(ctx, restored.SessionID, "task-user", "new", "", QueuedByUser, false, nil) + require.NoError(t, err) + require.False(t, appended) + require.NotEqual(t, restored.ID, inserted.ID) + entries, _, err := svc.SnapshotSession(ctx, restored.SessionID) + require.NoError(t, err) + require.Len(t, entries, 3) + require.Equal(t, restored.ID, entries[0].ID) + require.Equal(t, tail.ID, entries[1].ID) + require.Equal(t, inserted.ID, entries[2].ID) + }) + }) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/edit_lease_purge_race_test.go b/apps/backend/internal/orchestrator/messagequeue/edit_lease_purge_race_test.go new file mode 100644 index 0000000000..96e82e032f --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/edit_lease_purge_race_test.go @@ -0,0 +1,128 @@ +package messagequeue + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type blockingPurgeRepository struct { + Repository + purgeStarted chan struct{} + purgeRelease chan struct{} + listStarted chan struct{} + insertStarted chan struct{} + insertRelease chan struct{} + listOnce sync.Once +} + +func (r *blockingPurgeRepository) PurgeTask(ctx context.Context, taskID string) (int, error) { + close(r.purgeStarted) + <-r.purgeRelease + return r.Repository.PurgeTask(ctx, taskID) +} + +func (r *blockingPurgeRepository) ListBySession(ctx context.Context, sessionID string) ([]QueuedMessage, error) { + r.listOnce.Do(func() { close(r.listStarted) }) + return r.Repository.ListBySession(ctx, sessionID) +} + +func (r *blockingPurgeRepository) Insert(ctx context.Context, msg *QueuedMessage, maxPerSession int) error { + r.blockInsert() + return r.Repository.Insert(ctx, msg, maxPerSession) +} + +func (r *blockingPurgeRepository) InsertForSession(ctx context.Context, identity QueueSessionIdentity, msg *QueuedMessage, maxPerSession int) error { + r.blockInsert() + return r.Repository.InsertForSession(ctx, identity, msg, maxPerSession) +} + +func (r *blockingPurgeRepository) blockInsert() { + if r.insertStarted != nil { + close(r.insertStarted) + <-r.insertRelease + } +} + +func TestPurgeTaskWaitsForInFlightQueueAdmission(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + blocking := &blockingPurgeRepository{ + Repository: svc.repo, + purgeStarted: make(chan struct{}), + purgeRelease: make(chan struct{}), + insertStarted: make(chan struct{}), + insertRelease: make(chan struct{}), + } + svc.repo = blocking + + queueDone := make(chan error, 1) + go func() { + _, queueErr := svc.QueueMessage(ctx, "session-purge-admission", "task-purge-admission", "queued", "", QueuedByUser, false, nil) + queueDone <- queueErr + }() + <-blocking.insertStarted + + purgeDone := make(chan error, 1) + go func() { + _, purgeErr := svc.PurgeTask(ctx, "task-purge-admission") + purgeDone <- purgeErr + }() + + select { + case <-blocking.purgeStarted: + t.Fatal("PurgeTask reached the repository while queue admission was in flight") + case <-time.After(100 * time.Millisecond): + } + + close(blocking.insertRelease) + require.NoError(t, <-queueDone) + select { + case <-blocking.purgeStarted: + case <-time.After(time.Second): + t.Fatal("PurgeTask did not proceed after queue admission completed") + } + close(blocking.purgeRelease) + require.NoError(t, <-purgeDone) +} + +func TestPurgeTaskBlocksBeginEditUntilLeaseInvalidation(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-purge-race", "task-purge-race", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + + blocking := &blockingPurgeRepository{ + Repository: svc.repo, + purgeStarted: make(chan struct{}), + purgeRelease: make(chan struct{}), + listStarted: make(chan struct{}), + } + svc.repo = blocking + + purgeDone := make(chan error, 1) + go func() { + _, purgeErr := svc.PurgeTask(ctx, entry.TaskID) + purgeDone <- purgeErr + }() + <-blocking.purgeStarted + + beginDone := make(chan error, 1) + go func() { + _, beginErr := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + beginDone <- beginErr + }() + + select { + case <-blocking.listStarted: + t.Fatal("BeginEdit read the queue while PurgeTask was in progress") + case <-time.After(100 * time.Millisecond): + } + + close(blocking.purgeRelease) + require.NoError(t, <-purgeDone) + require.ErrorIs(t, <-beginDone, ErrEditLeaseNotFound) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/edit_lease_test.go b/apps/backend/internal/orchestrator/messagequeue/edit_lease_test.go new file mode 100644 index 0000000000..7622f10967 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/edit_lease_test.go @@ -0,0 +1,285 @@ +package messagequeue + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestEditLeaseProtectsTargetWithoutPausingQueuePolicy(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + first, err := svc.QueueMessage(ctx, "session-lease", "task", "first", "", "user", false, nil) + require.NoError(t, err) + second, err := svc.QueueMessage(ctx, "session-lease", "task", "second", "", "user", false, nil) + require.NoError(t, err) + + lease, err := svc.BeginEdit(ctx, second.SessionID, second.ID, "connection-a") + require.NoError(t, err) + require.Equal(t, int64(0), lease.TargetRevision) + + got, ok, autoRun := svc.ReserveQueuedWithAutoRun(ctx, first.SessionID) + require.True(t, autoRun) + require.True(t, ok) + require.Equal(t, first.ID, got.ID) + got, ok, autoRun = svc.ReserveQueuedWithAutoRun(ctx, second.SessionID) + require.True(t, autoRun) + require.False(t, ok) + require.Nil(t, got) + + _, err = svc.BeginEdit(ctx, second.SessionID, second.ID, "connection-b") + require.ErrorIs(t, err, ErrEditConflict) + + revision, err := svc.UpdateMessageWithLease(ctx, second.SessionID, second.ID, lease.LeaseID, + "operation-1", "connection-a", lease.TargetRevision, "edited", nil, nil) + require.NoError(t, err) + require.Equal(t, int64(1), revision) + + revision, err = svc.UpdateMessageWithLease(ctx, second.SessionID, second.ID, lease.LeaseID, + "operation-1", "connection-a", lease.TargetRevision, "edited", nil, nil) + require.NoError(t, err) + require.Equal(t, int64(1), revision) + + _, err = svc.UpdateMessageWithLease(ctx, second.SessionID, second.ID, lease.LeaseID, + "operation-2", "connection-a", lease.TargetRevision, "stale", nil, nil) + require.ErrorIs(t, err, ErrEditRevisionConflict) + + require.NoError(t, svc.EndEdit(ctx, second.SessionID, second.ID, lease.LeaseID, "connection-a")) + got, ok, _ = svc.ReserveQueuedWithAutoRun(ctx, second.SessionID) + require.True(t, ok) + require.Equal(t, second.ID, got.ID) +} + +func TestEditLeaseRejectsOperationHashMismatchWithoutMutation(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-operation", "task", "before", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + revision, err := svc.UpdateMessageWithLease( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation-1", "connection", + lease.TargetRevision, "after", nil, nil, + ) + require.NoError(t, err) + require.Equal(t, int64(1), revision) + + prepared := false + _, err = svc.UpdateMessageWithLeaseAfterValidation( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation-1", "connection", + lease.TargetRevision, "different", []MessageAttachment{{Type: "image", Data: "new"}}, + nil, + func(context.Context) error { + prepared = true + return nil + }, + nil, + ) + require.ErrorIs(t, err, ErrEditRevisionConflict) + require.False(t, prepared) + + status := svc.GetStatus(ctx, entry.SessionID) + require.Len(t, status.Entries, 1) + require.Equal(t, "after", status.Entries[0].Content) + require.Equal(t, int64(1), svc.editRevisions[svc.editLeaseKey(entry.SessionID, entry.ID)]) +} + +func TestReserveQueuedWithAutoRunReportsPolicyWhenEditBlocksHead(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-policy", "task", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + require.NoError(t, svc.SetAutoRun(ctx, entry.SessionID, false)) + _, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + got, ok, autoRun := svc.ReserveQueuedWithAutoRun(ctx, entry.SessionID) + require.False(t, autoRun) + require.False(t, ok) + require.Nil(t, got) +} + +func TestEditLeaseRejectsForeignReleaseAndUpdate(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-foreign", "task", "body", "", "user", false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + err = svc.EndEdit(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection-b") + require.ErrorIs(t, err, ErrEditLeaseNotFound) + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection-b", 0, "tampered", nil, nil) + require.ErrorIs(t, err, ErrEditLeaseNotFound) + require.False(t, errors.Is(err, ErrEditRevisionConflict)) +} +func TestEditLeaseEndRejectsExpiredLease(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-expiry", "task", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + svc.editLeases[svc.editLeaseKey(entry.SessionID, entry.ID)].ExpiresAt = time.Now().UTC().Add(-time.Second) + + err = svc.EndEdit(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection-a") + require.ErrorIs(t, err, ErrEditLeaseNotFound) + + _, ok := svc.TakeQueued(ctx, entry.SessionID) + require.True(t, ok, "an expired lease must not continue blocking delivery") +} + +func TestEditLeaseBlocksTargetedDrains(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-target", "task", "body", "", "user", false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + got, ok := svc.TakeQueued(ctx, entry.SessionID) + require.False(t, ok) + require.Nil(t, got) + _, err = svc.ClaimSendNow(ctx, entry.SessionID, []QueuedMessage{*entry}) + require.ErrorIs(t, err, ErrEditConflict) +} +func TestClaimSendNowClearsOrdinaryEditState(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-send-now", "task", "before", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection", lease.TargetRevision, "after", nil, nil) + require.NoError(t, err) + + key := svc.editLeaseKey(entry.SessionID, entry.ID) + require.Contains(t, svc.editRevisions, key) + require.Contains(t, svc.editRevisionTaskIDs, key) + svc.editLeases[key].ExpiresAt = time.Now().UTC().Add(-time.Second) + current, err := svc.GetEntry(ctx, entry.SessionID, entry.ID) + require.NoError(t, err) + claim, err := svc.ClaimSendNow(ctx, entry.SessionID, []QueuedMessage{*current}) + require.NoError(t, err) + require.Len(t, claim.Sources, 1) + require.NotContains(t, svc.editLeases, key) + require.NotContains(t, svc.editRevisions, key) + require.NotContains(t, svc.editRevisionTaskIDs, key) +} + +func TestEditLeaseReleaseForDisconnectedConnectionUnblocksTarget(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-disconnect", "task", "body", "", QueuedByUser, false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + require.Equal(t, 1, svc.ReleaseEditLeasesForConnection("connection-a")) + + got, ok := svc.TakeQueued(ctx, entry.SessionID) + require.True(t, ok) + require.Equal(t, entry.ID, got.ID) +} + +func TestTakeQueuedIfAutoRunDoesNotDeleteDurableLifecycleEntry(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, _, accepted, err := svc.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-lifecycle-auto-run", "task", "lifecycle", "", QueuedByWorkflow, + false, nil, map[string]interface{}{"origin": "github_pr_automation"}, "lifecycle-key", true, + ) + require.NoError(t, err) + require.True(t, accepted) + + got, ok := svc.TakeQueuedIfAutoRun(ctx, entry.SessionID) + require.False(t, ok) + require.Nil(t, got) + + status := svc.GetStatus(ctx, entry.SessionID) + require.Len(t, status.Entries, 1) + require.Equal(t, entry.ID, status.Entries[0].ID) +} + +func TestLeasedUpdateRollsBackPreparedStateOnPreparationError(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-prepare", "task", "before", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + claimed, rolledBack := false, false + _, err = svc.UpdateMessageWithLeaseAfterValidation( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation-1", "connection-a", + lease.TargetRevision, "after", nil, nil, + func(context.Context) error { + claimed = true + return errors.New("attachment claim failed") + }, + func(context.Context) error { + rolledBack = true + return nil + }, + ) + require.Error(t, err) + require.True(t, claimed) + require.True(t, rolledBack) + + status := svc.GetStatus(ctx, entry.SessionID) + require.Equal(t, "before", status.Entries[0].Content) +} + +func TestQueueCoalesceReplacementInvalidatesEditLease(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, replaced, err := svc.QueueMessageWithCoalesceKey( + ctx, "session-lease-replace", "task", "before", "", QueuedByUser, + false, nil, nil, "coalesce-key", true, + ) + require.NoError(t, err) + require.False(t, replaced) + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + replacement, replaced, err := svc.QueueMessageWithCoalesceKey( + ctx, entry.SessionID, "task", "replacement", "", QueuedByUser, + false, nil, nil, "coalesce-key", true, + ) + require.NoError(t, err) + require.True(t, replaced) + require.Equal(t, entry.ID, replacement.ID) + + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection-a", lease.TargetRevision, "stale", nil, nil) + require.ErrorIs(t, err, ErrEditLeaseNotFound) +} +func TestEndEditAfterSaveReportsOnlyFinalizedUpdates(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-lease-save", "task", "before", "", QueuedByUser, false, nil) + require.NoError(t, err) + + lease, err := svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + saved, err := svc.EndEditAfterSave(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection") + require.NoError(t, err) + require.False(t, saved) + + lease, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + _, err = svc.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, + "operation-1", "connection", lease.TargetRevision, "after", nil, nil) + require.NoError(t, err) + + saved, err = svc.EndEditAfterSave(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection") + require.NoError(t, err) + require.True(t, saved) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/purge_pending_move_test.go b/apps/backend/internal/orchestrator/messagequeue/purge_pending_move_test.go new file mode 100644 index 0000000000..62ce9df3a4 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/purge_pending_move_test.go @@ -0,0 +1,54 @@ +package messagequeue + +import ( + "context" + "testing" + + _ "github.com/mattn/go-sqlite3" + "github.com/stretchr/testify/require" +) + +func TestRepositories_PurgeTaskRemovesPendingMove(t *testing.T) { + factories := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + for _, factory := range factories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + ctx := context.Background() + require.NoError(t, repo.SetPendingMove(ctx, "session-purge", &PendingMove{ + TaskID: "task-purge", + })) + + _, err := repo.PurgeTask(ctx, "task-purge") + require.NoError(t, err) + + move, err := repo.GetPendingMove(ctx, "session-purge") + require.NoError(t, err) + require.Nil(t, move) + }) + } +} + +func TestSQLitePurgeTaskInTransactionRemovesPendingMove(t *testing.T) { + repo := newTestSQLiteRepo(t).(*sqliteRepository) + ctx := context.Background() + require.NoError(t, repo.SetPendingMove(ctx, "session-purge-tx", &PendingMove{ + TaskID: "task-purge-tx", + })) + + tx, err := repo.db.BeginTxx(ctx, nil) + require.NoError(t, err) + removed, err := PurgeTaskInTransaction(ctx, tx, repo.db, "task-purge-tx", []string{"session-purge-tx"}) + require.NoError(t, err) + require.Zero(t, removed) + require.NoError(t, tx.Commit()) + + move, err := repo.GetPendingMove(ctx, "session-purge-tx") + require.NoError(t, err) + require.Nil(t, move) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository.go b/apps/backend/internal/orchestrator/messagequeue/repository.go index afd2389ddc..d1c069eda1 100644 --- a/apps/backend/internal/orchestrator/messagequeue/repository.go +++ b/apps/backend/internal/orchestrator/messagequeue/repository.go @@ -74,6 +74,10 @@ type Repository interface { // task. Lifecycle insertion verifies this generation atomically. LifecycleGeneration(ctx context.Context, taskID string) (int64, error) + // SessionGeneration returns the current destructive-mutation generation for + // a session. Send Now restores use it to reject stale FIFO reservations. + SessionGeneration(ctx context.Context, sessionID string) (int64, error) + // PurgeTask is backend-only cleanup. It removes all task rows, including // reserved server-owned lifecycle rows, and advances its generation. PurgeTask(ctx context.Context, taskID string) (int, error) @@ -102,7 +106,8 @@ type Repository interface { // ReserveHead returns the lowest-position entry. Ordinary entries are // atomically deleted, matching TakeHead. Durable lifecycle entries remain - // stored until AcknowledgeByID is called after executor acceptance. + // stored until AcknowledgeReserved receives the exact reservation returned + // by this call after executor acceptance. ReserveHead(ctx context.Context, sessionID string) (*QueuedMessage, error) // GetAutoRun returns the durable per-session automatic-drain policy. Missing @@ -138,6 +143,9 @@ type Repository interface { identity QueueSessionIdentity, ) (*QueuedMessage, bool, error) + // AcknowledgeReserved removes only the exact lifecycle reservation carried + // by msg. A stale delivery attempt cannot remove a newer retry. + AcknowledgeReserved(ctx context.Context, msg *QueuedMessage) error // AcknowledgeByID is an internal dispatch operation that removes a reserved // entry regardless of its server-owned queued_by identity. AcknowledgeByID(ctx context.Context, sessionID, entryID string) error @@ -249,6 +257,11 @@ type Repository interface { // require identity to remain current: stale workers use it after replacement. DiscardLifecycleReservation(ctx context.Context, identity QueueSessionIdentity, entryID string) error + // PurgeSession removes every queue row for a deleted session, including + // durable lifecycle rows reserved in flight, and its pending workflow move. + // Returns the exact number of queue rows removed. + PurgeSession(ctx context.Context, sessionID string) (int, error) + // TransferSession moves all entries (and any pending move) from oldSessionID // to newSessionID. Used on workflow session switches. TransferSession(ctx context.Context, oldSessionID, newSessionID string) error diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_attachment_cleanup.go b/apps/backend/internal/orchestrator/messagequeue/repository_attachment_cleanup.go new file mode 100644 index 0000000000..9b8ec0e751 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_attachment_cleanup.go @@ -0,0 +1,291 @@ +package messagequeue + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "github.com/jmoiron/sqlx" + internaldb "github.com/kandev/kandev/internal/db" + "sort" + "time" +) + +const attachmentCleanupSchema = ` + CREATE TABLE IF NOT EXISTS queue_attachment_cleanups ( + session_id TEXT NOT NULL, + current_session_id TEXT NOT NULL DEFAULT '', + entry_id TEXT NOT NULL, + operation_id TEXT NOT NULL DEFAULT '', + task_id TEXT NOT NULL, + owner_id TEXT NOT NULL DEFAULT '', + lease_id TEXT NOT NULL DEFAULT '', + remove_entry INTEGER NOT NULL DEFAULT 0, + claim_pending INTEGER NOT NULL DEFAULT 0, + entry_fingerprint TEXT NOT NULL DEFAULT '', + attachments_json TEXT NOT NULL DEFAULT '[]', + created_at TIMESTAMP NOT NULL, + PRIMARY KEY (session_id, entry_id, operation_id) + ) +` + +func (r *sqliteRepository) ensureAttachmentCleanupSchema(ctx context.Context) error { + if _, err := r.db.ExecContext(ctx, attachmentCleanupSchema); err != nil { + return fmt.Errorf("ensure attachment cleanup schema: %w", err) + } + if _, err := r.db.ExecContext(ctx, `ALTER TABLE queue_attachment_cleanups ADD COLUMN owner_id TEXT NOT NULL DEFAULT ''`); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add attachment cleanup owner: %w", err) + } + if _, err := r.db.ExecContext(ctx, `ALTER TABLE queue_attachment_cleanups ADD COLUMN remove_entry INTEGER NOT NULL DEFAULT 0`); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add attachment cleanup remove-entry state: %w", err) + } + if _, err := r.db.ExecContext(ctx, `ALTER TABLE queue_attachment_cleanups ADD COLUMN claim_pending INTEGER NOT NULL DEFAULT 0`); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add attachment cleanup claim-pending state: %w", err) + } + if _, err := r.db.ExecContext(ctx, `ALTER TABLE queue_attachment_cleanups ADD COLUMN entry_fingerprint TEXT NOT NULL DEFAULT ''`); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add attachment cleanup entry fingerprint: %w", err) + } + if _, err := r.db.ExecContext(ctx, `ALTER TABLE queue_attachment_cleanups ADD COLUMN current_session_id TEXT NOT NULL DEFAULT ''`); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add attachment cleanup current session: %w", err) + } + if _, err := r.db.ExecContext(ctx, ` + UPDATE queue_attachment_cleanups + SET current_session_id = session_id + WHERE current_session_id = '' + `); err != nil { + return fmt.Errorf("backfill attachment cleanup current session: %w", err) + } + return nil +} + +func (r *sqliteRepository) UpsertAttachmentCleanup(ctx context.Context, cleanup AttachmentCleanup) error { + if err := r.ensureAttachmentCleanupSchema(ctx); err != nil { + return err + } + attachmentsJSON, err := json.Marshal(cleanup.Attachments) + if err != nil { + return fmt.Errorf("marshal attachment cleanup: %w", err) + } + if cleanup.CreatedAt.IsZero() { + cleanup.CreatedAt = time.Now().UTC() + } + if cleanup.CurrentSessionID == "" { + cleanup.CurrentSessionID = cleanup.SessionID + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin attachment cleanup upsert: %w", err) + } + defer func() { _ = tx.Rollback() }() + cleanup.CurrentSessionID, err = r.resolveAttachmentCleanupSessionTx(ctx, tx, cleanup) + if err != nil { + return err + } + if _, err := tx.ExecContext(ctx, tx.Rebind(` + INSERT INTO queue_attachment_cleanups + (session_id, current_session_id, entry_id, operation_id, task_id, owner_id, lease_id, + remove_entry, claim_pending, entry_fingerprint, attachments_json, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(session_id, entry_id, operation_id) DO UPDATE SET + task_id = excluded.task_id, + owner_id = excluded.owner_id, + lease_id = excluded.lease_id, + remove_entry = excluded.remove_entry, + claim_pending = excluded.claim_pending, + entry_fingerprint = excluded.entry_fingerprint, + attachments_json = excluded.attachments_json + `), cleanup.SessionID, cleanup.CurrentSessionID, cleanup.EntryID, cleanup.OperationID, cleanup.TaskID, + cleanup.OwnerID, cleanup.LeaseID, boolToInt(cleanup.RemoveEntry), + boolToInt(cleanup.ClaimPending), cleanup.EntryFingerprint, + string(attachmentsJSON), cleanup.CreatedAt); err != nil { + return fmt.Errorf("upsert attachment cleanup: %w", err) + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit attachment cleanup upsert: %w", err) + } + return nil +} + +func (r *sqliteRepository) resolveAttachmentCleanupSessionTx( + ctx context.Context, + tx *sqlx.Tx, + cleanup AttachmentCleanup, +) (string, error) { + entrySessionID := cleanup.CurrentSessionID + queuedSessionID, err := r.queuedAttachmentCleanupSessionTx(ctx, tx, cleanup) + if err != nil { + return "", err + } + if queuedSessionID != "" { + entrySessionID = queuedSessionID + } + sessionIDs := make([]string, 0, 2) + seenSessionIDs := make(map[string]struct{}, 2) + for _, sessionID := range []string{cleanup.CurrentSessionID, entrySessionID} { + if sessionID == "" { + continue + } + if _, seen := seenSessionIDs[sessionID]; seen { + continue + } + seenSessionIDs[sessionID] = struct{}{} + sessionIDs = append(sessionIDs, sessionID) + } + sort.Strings(sessionIDs) + for _, sessionID := range sessionIDs { + if err := r.lockSessionTxUnfenced(ctx, tx, sessionID); err != nil { + return "", err + } + } + queuedSessionID, err = r.queuedAttachmentCleanupSessionTx(ctx, tx, cleanup) + if err != nil { + return "", err + } + if queuedSessionID != "" { + entrySessionID = queuedSessionID + } + return ResolveSessionTransferInTransaction(ctx, tx, r.db, entrySessionID) +} + +func (r *sqliteRepository) queuedAttachmentCleanupSessionTx( + ctx context.Context, + tx *sqlx.Tx, + cleanup AttachmentCleanup, +) (string, error) { + var queuedSessionID string + err := tx.GetContext(ctx, &queuedSessionID, tx.Rebind(` + SELECT session_id + FROM queued_messages + WHERE id = ? AND task_id = ? + `), cleanup.EntryID, cleanup.TaskID) + if errors.Is(err, sql.ErrNoRows) { + return "", nil + } + if err != nil { + return "", fmt.Errorf("resolve queued attachment cleanup entry: %w", err) + } + return queuedSessionID, nil +} +func (r *sqliteRepository) DeleteAttachmentCleanup( + ctx context.Context, + sessionID, entryID, operationID string, +) error { + if err := r.ensureAttachmentCleanupSchema(ctx); err != nil { + return err + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin attachment cleanup delete: %w", err) + } + defer func() { _ = tx.Rollback() }() + var currentSessionID string + err = tx.GetContext(ctx, ¤tSessionID, tx.Rebind(` + SELECT current_session_id FROM queue_attachment_cleanups + WHERE session_id = ? AND entry_id = ? AND operation_id = ? + `), sessionID, entryID, operationID) + if errors.Is(err, sql.ErrNoRows) { + return nil + } + if err != nil { + return fmt.Errorf("locate attachment cleanup for delete: %w", err) + } + if err := r.lockSessionTx(ctx, tx, currentSessionID); err != nil { + return err + } + if _, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM queue_attachment_cleanups + WHERE session_id = ? AND entry_id = ? AND operation_id = ? + `), sessionID, entryID, operationID); err != nil { + return fmt.Errorf("delete attachment cleanup: %w", err) + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit attachment cleanup delete: %w", err) + } + return nil +} + +func (r *sqliteRepository) ListAttachmentCleanups(ctx context.Context) ([]AttachmentCleanup, error) { + if err := r.ensureAttachmentCleanupSchema(ctx); err != nil { + return nil, err + } + rows, err := r.db.QueryxContext(ctx, ` + SELECT session_id, current_session_id, entry_id, operation_id, task_id, owner_id, lease_id, + remove_entry, claim_pending, entry_fingerprint, attachments_json, created_at + FROM queue_attachment_cleanups + ORDER BY created_at, session_id, entry_id, operation_id + `) + if err != nil { + return nil, fmt.Errorf("list attachment cleanups: %w", err) + } + defer func() { _ = rows.Close() }() + var cleanups []AttachmentCleanup + for rows.Next() { + var cleanup AttachmentCleanup + var attachmentsJSON string + if err := rows.Scan( + &cleanup.SessionID, + &cleanup.CurrentSessionID, + &cleanup.EntryID, + &cleanup.OperationID, + &cleanup.TaskID, + &cleanup.OwnerID, + &cleanup.LeaseID, + &cleanup.RemoveEntry, + &cleanup.ClaimPending, + &cleanup.EntryFingerprint, + &attachmentsJSON, + &cleanup.CreatedAt, + ); err != nil { + return nil, fmt.Errorf("scan attachment cleanup: %w", err) + } + if err := json.Unmarshal([]byte(attachmentsJSON), &cleanup.Attachments); err != nil { + return nil, fmt.Errorf("unmarshal attachment cleanup: %w", err) + } + cleanups = append(cleanups, cleanup) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate attachment cleanups: %w", err) + } + return cleanups, nil +} + +func (r *sqliteRepository) GetAttachmentCleanup( + ctx context.Context, + sessionID, entryID, operationID string, +) (*AttachmentCleanup, error) { + if err := r.ensureAttachmentCleanupSchema(ctx); err != nil { + return nil, err + } + var cleanup AttachmentCleanup + var attachmentsJSON string + err := r.db.QueryRowxContext(ctx, r.db.Rebind(` + SELECT session_id, current_session_id, entry_id, operation_id, task_id, owner_id, lease_id, + remove_entry, claim_pending, entry_fingerprint, attachments_json, created_at + FROM queue_attachment_cleanups + WHERE session_id = ? AND entry_id = ? AND operation_id = ? + `), sessionID, entryID, operationID).Scan( + &cleanup.SessionID, + &cleanup.CurrentSessionID, + &cleanup.EntryID, + &cleanup.OperationID, + &cleanup.TaskID, + &cleanup.OwnerID, + &cleanup.LeaseID, + &cleanup.RemoveEntry, + &cleanup.ClaimPending, + &cleanup.EntryFingerprint, + &attachmentsJSON, + &cleanup.CreatedAt, + ) + if err == sql.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("get attachment cleanup: %w", err) + } + if err := json.Unmarshal([]byte(attachmentsJSON), &cleanup.Attachments); err != nil { + return nil, fmt.Errorf("unmarshal attachment cleanup: %w", err) + } + return &cleanup, nil +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_dispatch_recovery.go b/apps/backend/internal/orchestrator/messagequeue/repository_dispatch_recovery.go new file mode 100644 index 0000000000..18cbcba51c --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_dispatch_recovery.go @@ -0,0 +1,318 @@ +package messagequeue + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "github.com/google/uuid" + "time" + + "github.com/jmoiron/sqlx" + internaldb "github.com/kandev/kandev/internal/db" +) + +// PendingQueueDispatch records an ordinary queue row removed for at-least-once +// delivery. Accepted selects acknowledgement rather than restoration during +// startup. A crash after agent acceptance but before Accepted commits can +// deliver the same prompt again; the agent protocol has no idempotency key. +type PendingQueueDispatch struct { + Message QueuedMessage + Accepted bool +} + +const queueDispatchRecoverySchema = ` + CREATE TABLE IF NOT EXISTS queue_dispatch_claims ( + entry_id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + attempt_id TEXT NOT NULL, + message_json TEXT NOT NULL, + accepted INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMP NOT NULL + ) +` + +func (r *sqliteRepository) ensureQueueDispatchRecoverySchema(ctx context.Context) error { + if _, err := r.db.ExecContext(ctx, queueDispatchRecoverySchema); err != nil { + return fmt.Errorf("ensure queue dispatch recovery schema: %w", err) + } + if _, err := r.db.ExecContext(ctx, ` + ALTER TABLE queue_dispatch_claims ADD COLUMN attempt_id TEXT NOT NULL DEFAULT '' + `); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add queue dispatch attempt id: %w", err) + } + rows, err := r.db.QueryxContext(ctx, `SELECT entry_id FROM queue_dispatch_claims WHERE attempt_id = ''`) + if err != nil { + return fmt.Errorf("list queue dispatches without attempt ids: %w", err) + } + var entryIDs []string + for rows.Next() { + var entryID string + if err := rows.Scan(&entryID); err != nil { + _ = rows.Close() + return fmt.Errorf("scan queue dispatch without attempt id: %w", err) + } + entryIDs = append(entryIDs, entryID) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return fmt.Errorf("iterate queue dispatches without attempt ids: %w", err) + } + if err := rows.Close(); err != nil { + return fmt.Errorf("close queue dispatch attempt migration rows: %w", err) + } + for _, entryID := range entryIDs { + if _, err := r.db.ExecContext(ctx, r.db.Rebind(` + UPDATE queue_dispatch_claims SET attempt_id = ? + WHERE entry_id = ? AND attempt_id = '' + `), uuid.NewString(), entryID); err != nil { + return fmt.Errorf("backfill queue dispatch attempt id: %w", err) + } + } + return nil +} + +func (r *sqliteRepository) persistQueueDispatchClaimTx( + ctx context.Context, + tx *sqlx.Tx, + msg *QueuedMessage, +) error { + msg.dispatchAttemptID = uuid.NewString() + messageJSON, err := json.Marshal(msg) + if err != nil { + return fmt.Errorf("marshal queue dispatch claim: %w", err) + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + INSERT INTO queue_dispatch_claims (entry_id, session_id, attempt_id, message_json, created_at) + VALUES (?, ?, ?, ?, ?) + `), msg.ID, msg.SessionID, msg.dispatchAttemptID, string(messageJSON), time.Now().UTC()); err != nil { + return fmt.Errorf("persist queue dispatch claim: %w", err) + } + return nil +} + +func (r *sqliteRepository) ListPendingQueueDispatches(ctx context.Context) ([]PendingQueueDispatch, error) { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return nil, err + } + rows, err := r.db.QueryxContext(ctx, ` + SELECT attempt_id, message_json, accepted + FROM queue_dispatch_claims ORDER BY created_at, entry_id + `) + if err != nil { + return nil, fmt.Errorf("list pending queue dispatch: %w", err) + } + defer func() { _ = rows.Close() }() + var pending []PendingQueueDispatch + for rows.Next() { + var attemptID, messageJSON string + var accepted int + if err := rows.Scan(&attemptID, &messageJSON, &accepted); err != nil { + return nil, fmt.Errorf("scan pending queue dispatch: %w", err) + } + var msg QueuedMessage + if err := json.Unmarshal([]byte(messageJSON), &msg); err != nil { + return nil, fmt.Errorf("unmarshal pending queue dispatch: %w", err) + } + msg.dispatchAttemptID = attemptID + msg.reservationGenerationsCaptured = true + pending = append(pending, PendingQueueDispatch{Message: msg, Accepted: accepted != 0}) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate pending queue dispatch: %w", err) + } + return pending, nil +} + +func (r *sqliteRepository) transferPendingQueueDispatchesTx( + ctx context.Context, + tx *sqlx.Tx, + oldSessionID, newSessionID string, +) error { + rows, err := tx.QueryxContext(ctx, r.db.Rebind(` + SELECT entry_id, message_json FROM queue_dispatch_claims WHERE session_id = ? + `), oldSessionID) + if err != nil { + return fmt.Errorf("list transferred queue dispatch claims: %w", err) + } + type claimUpdate struct { + entryID string + attemptID string + messageJSON string + } + var updates []claimUpdate + for rows.Next() { + var entryID, messageJSON string + if err := rows.Scan(&entryID, &messageJSON); err != nil { + _ = rows.Close() + return fmt.Errorf("scan transferred queue dispatch claim: %w", err) + } + var msg QueuedMessage + if err := json.Unmarshal([]byte(messageJSON), &msg); err != nil { + _ = rows.Close() + return fmt.Errorf("unmarshal transferred queue dispatch claim: %w", err) + } + msg.SessionID = newSessionID + attemptID := uuid.NewString() + updatedJSON, err := json.Marshal(msg) + if err != nil { + _ = rows.Close() + return fmt.Errorf("marshal transferred queue dispatch claim: %w", err) + } + updates = append(updates, claimUpdate{ + entryID: entryID, attemptID: attemptID, messageJSON: string(updatedJSON), + }) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return fmt.Errorf("iterate transferred queue dispatch claims: %w", err) + } + if err := rows.Close(); err != nil { + return fmt.Errorf("close transferred queue dispatch claims: %w", err) + } + for _, update := range updates { + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + UPDATE queue_dispatch_claims SET session_id = ?, attempt_id = ?, message_json = ? + WHERE entry_id = ? AND session_id = ? + `), newSessionID, update.attemptID, update.messageJSON, update.entryID, oldSessionID); err != nil { + return fmt.Errorf("transfer queue dispatch claim: %w", err) + } + } + return nil +} +func pendingQueueDispatchesForTaskTx( + ctx context.Context, + tx *sqlx.Tx, + taskID string, +) ([]string, []string, error) { + rows, err := tx.QueryxContext(ctx, `SELECT entry_id, session_id, message_json FROM queue_dispatch_claims`) + if err != nil { + return nil, nil, fmt.Errorf("list task queue dispatch claims: %w", err) + } + defer func() { _ = rows.Close() }() + var entryIDs []string + var sessions []string + for rows.Next() { + var entryID, sessionID, messageJSON string + if err := rows.Scan(&entryID, &sessionID, &messageJSON); err != nil { + return nil, nil, fmt.Errorf("scan task queue dispatch claim: %w", err) + } + var msg QueuedMessage + if err := json.Unmarshal([]byte(messageJSON), &msg); err != nil { + return nil, nil, fmt.Errorf("unmarshal task queue dispatch claim: %w", err) + } + if msg.TaskID == taskID { + entryIDs = append(entryIDs, entryID) + sessions = append(sessions, sessionID) + } + } + if err := rows.Err(); err != nil { + return nil, nil, fmt.Errorf("iterate task queue dispatch claims: %w", err) + } + return entryIDs, sessions, nil +} + +func (r *sqliteRepository) validatePendingQueueDispatchTx( + ctx context.Context, + tx *sqlx.Tx, + msg *QueuedMessage, +) error { + var storedSessionID, storedAttemptID string + var accepted int + err := tx.QueryRowxContext(ctx, r.db.Rebind(` + SELECT session_id, attempt_id, accepted FROM queue_dispatch_claims WHERE entry_id = ? + `), msg.ID).Scan(&storedSessionID, &storedAttemptID, &accepted) + if errors.Is(err, sql.ErrNoRows) { + return ErrQueueDispatchClaimChanged + } + if err != nil { + return fmt.Errorf("read queue dispatch claim: %w", err) + } + if storedSessionID != msg.SessionID || storedAttemptID != msg.dispatchAttemptID || accepted != 0 { + return ErrQueueDispatchClaimChanged + } + return nil +} + +func (r *sqliteRepository) deletePendingQueueDispatchTx( + ctx context.Context, + tx *sqlx.Tx, + msg *QueuedMessage, +) error { + result, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queue_dispatch_claims + WHERE session_id = ? AND entry_id = ? AND attempt_id = ? + `), msg.SessionID, msg.ID, msg.dispatchAttemptID) + if err != nil { + return fmt.Errorf("delete queue dispatch claim: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("delete queue dispatch claim rows affected: %w", err) + } + if affected != 1 { + return ErrQueueDispatchClaimChanged + } + return nil +} + +func (r *sqliteRepository) deletePendingQueueDispatchesBySessionTx( + ctx context.Context, + tx *sqlx.Tx, + sessionID string, +) error { + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queue_dispatch_claims WHERE session_id = ? + `), sessionID); err != nil { + return fmt.Errorf("delete session queue dispatch claims: %w", err) + } + return nil +} + +func (r *sqliteRepository) MarkPendingQueueDispatchAccepted( + ctx context.Context, + msg *QueuedMessage, +) error { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return err + } + tx, err := r.beginSessionMutationTx(ctx, msg.SessionID, "mark queue dispatch accepted") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + result, err := tx.ExecContext(ctx, r.db.Rebind(` + UPDATE queue_dispatch_claims SET accepted = 1 + WHERE entry_id = ? AND session_id = ? AND attempt_id = ? + `), msg.ID, msg.SessionID, msg.dispatchAttemptID) + if err != nil { + return fmt.Errorf("mark queue dispatch accepted: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("mark queue dispatch accepted rows affected: %w", err) + } + if affected != 1 { + return ErrQueueDispatchClaimChanged + } + return tx.Commit() +} + +func (r *sqliteRepository) DeletePendingQueueDispatch( + ctx context.Context, + msg *QueuedMessage, +) error { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return err + } + tx, err := r.beginSessionMutationTx(ctx, msg.SessionID, "delete pending queue dispatch") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + if err := r.deletePendingQueueDispatchTx(ctx, tx, msg); err != nil { + return err + } + return tx.Commit() +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_dispatch_session_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_dispatch_session_test.go new file mode 100644 index 0000000000..ffa5f6e757 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_dispatch_session_test.go @@ -0,0 +1,294 @@ +package messagequeue + +import ( + "context" + "encoding/json" + "errors" + "testing" +) + +func TestSQLiteTransferSessionMovesOrdinaryDispatchClaim(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + source := insertTestEntry(t, repo, "session-transfer-old", "task-transfer", "prompt", QueuedByUser, nil, nil) + reserved, _, err := repo.ReserveHeadIfAutoRun(ctx, source.SessionID) + if err != nil || reserved == nil { + t.Fatalf("reserve ordinary head = %#v, err=%v", reserved, err) + } + if err := repo.TransferSession(ctx, source.SessionID, "session-transfer-new"); err != nil { + t.Fatal(err) + } + pending, err := repo.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || pending[0].Message.SessionID != "session-transfer-new" { + t.Fatalf("dispatch claims after transfer = %#v", pending) + } + if err := repo.MarkPendingQueueDispatchAccepted(ctx, reserved); !errors.Is(err, ErrQueueDispatchClaimChanged) { + t.Fatalf("stale migrated claim acceptance = %v, want %v", err, ErrQueueDispatchClaimChanged) + } + if err := repo.DeletePendingQueueDispatch(ctx, reserved); !errors.Is(err, ErrQueueDispatchClaimChanged) { + t.Fatalf("stale migrated claim acknowledgement = %v, want %v", err, ErrQueueDispatchClaimChanged) + } + pending, err = repo.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || pending[0].Message.dispatchAttemptID == reserved.dispatchAttemptID { + t.Fatalf("transferred claim attempt was not rotated: reserved=%q pending=%#v", reserved.dispatchAttemptID, pending) + } + if err := repo.MarkPendingQueueDispatchAccepted(ctx, &pending[0].Message); err != nil { + t.Fatalf("mark current migrated claim: %v", err) + } +} + +func TestSQLiteReplaceSessionInvalidatesOrdinaryDispatchClaims(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + source := insertTestEntry(t, repo, "session-replace", "task-replace", "prompt", QueuedByUser, nil, nil) + reserved, _, err := repo.ReserveHeadIfAutoRun(ctx, source.SessionID) + if err != nil || reserved == nil { + t.Fatalf("reserve ordinary head = %#v, err=%v", reserved, err) + } + if err := repo.ReplaceSession(ctx, source.SessionID, nil, nil); err != nil { + t.Fatal(err) + } + pending, err := repo.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("dispatch claims after replacement = %#v, want empty", pending) + } +} + +func TestSQLiteStaleOrdinarySettlementCannotResurrectAfterSessionMutation(t *testing.T) { + tests := []struct { + name string + mutate func(context.Context, *Service, *QueuedMessage) error + settle func(context.Context, *Service, *QueuedMessage) error + }{ + { + name: "transfer before restore", + mutate: func(ctx context.Context, service *Service, msg *QueuedMessage) error { + return service.TransferSession(ctx, msg.SessionID, "session-transfer-destination") + }, + settle: func(ctx context.Context, service *Service, msg *QueuedMessage) error { + _, err := service.RestoreMessage(ctx, msg) + return err + }, + }, + { + name: "replace before requeue", + mutate: func(ctx context.Context, service *Service, msg *QueuedMessage) error { + return service.repo.ReplaceSession(ctx, msg.SessionID, nil, nil) + }, + settle: func(ctx context.Context, service *Service, msg *QueuedMessage) error { + return service.RequeueAtHead(ctx, msg) + }, + }, + { + name: "clear before restore", + mutate: func(ctx context.Context, service *Service, msg *QueuedMessage) error { + _, err := service.CancelAll(ctx, msg.SessionID) + return err + }, + settle: func(ctx context.Context, service *Service, msg *QueuedMessage) error { + _, err := service.RestoreMessage(ctx, msg) + return err + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + service := setupService(t) + service.repo = repo + source, err := service.QueueMessage( + ctx, "session-stale", "task-stale", "prompt", "", QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := service.ReserveQueued(ctx, source.SessionID) + if !ok || reserved.ID != source.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + if err := tc.mutate(ctx, service, reserved); err != nil { + t.Fatal(err) + } + if err := tc.settle(ctx, service, reserved); !errors.Is(err, ErrQueueDispatchClaimChanged) { + t.Fatalf("stale settlement error = %v, want %v", err, ErrQueueDispatchClaimChanged) + } + if entries, err := repo.ListBySession(ctx, source.SessionID); err != nil || len(entries) != 0 { + t.Fatalf("stale source queue = %#v, err=%v", entries, err) + } + }) + } +} + +func TestDestructiveTakeSettlementIsFencedAfterTransfer(t *testing.T) { + factories := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + takes := []struct { + name string + take func(context.Context, Repository, *QueuedMessage) (*QueuedMessage, error) + }{ + { + name: "head", + take: func(ctx context.Context, repo Repository, msg *QueuedMessage) (*QueuedMessage, error) { + return repo.TakeHead(ctx, msg.SessionID) + }, + }, + { + name: "by id", + take: func(ctx context.Context, repo Repository, msg *QueuedMessage) (*QueuedMessage, error) { + return repo.TakeByID(ctx, msg.SessionID, msg.ID) + }, + }, + } + for _, factory := range factories { + for _, take := range takes { + t.Run(factory.name+"/"+take.name, func(t *testing.T) { + ctx := context.Background() + repo := factory.new(t) + source := insertTestEntry(t, repo, "session-direct-old", "task-direct", "prompt", QueuedByUser, nil, nil) + reserved, err := take.take(ctx, repo, source) + if err != nil || reserved == nil { + t.Fatalf("take = %#v, err=%v", reserved, err) + } + if _, _, captured := reserved.ReservationGenerations(); !captured { + t.Fatal("destructive take did not capture settlement generation") + } + if err := repo.TransferSession(ctx, source.SessionID, "session-direct-new"); err != nil { + t.Fatal(err) + } + if err := repo.Restore(ctx, reserved, 0); !errors.Is(err, ErrQueueDispatchClaimChanged) { + t.Fatalf("stale direct-take restore = %v, want %v", err, ErrQueueDispatchClaimChanged) + } + }) + } + } +} + +func TestLifecycleAcknowledgementRequiresCurrentReservation(t *testing.T) { + factories := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + for _, factory := range factories { + t.Run(factory.name, func(t *testing.T) { + ctx := context.Background() + repo := factory.new(t) + source := insertTestEntry( + t, repo, "session-lifecycle-attempt", "task-lifecycle-attempt", "prompt", + QueuedByWorkflow, nil, map[string]interface{}{MetadataLifecycleDurable: true}, + ) + first, err := repo.ReserveHead(ctx, source.SessionID) + if err != nil || first == nil { + t.Fatalf("first reserve = %#v, err=%v", first, err) + } + second, err := repo.ReserveHead(ctx, source.SessionID) + if err != nil || second == nil { + t.Fatalf("second reserve = %#v, err=%v", second, err) + } + if first.lifecycleReservationID == "" || first.lifecycleReservationID == second.lifecycleReservationID { + t.Fatalf("lifecycle reservation ids = %q, %q", first.lifecycleReservationID, second.lifecycleReservationID) + } + if err := repo.AcknowledgeReserved(ctx, first); !errors.Is(err, ErrLifecycleReservationChanged) { + t.Fatalf("stale lifecycle acknowledgement = %v, want %v", err, ErrLifecycleReservationChanged) + } + if err := repo.AcknowledgeReserved(ctx, second); err != nil { + t.Fatalf("current lifecycle acknowledgement: %v", err) + } + }) + } +} + +func TestSQLiteQueueDispatchClaimMigratesLegacyAttemptID(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + _, err := repo.db.ExecContext(ctx, ` + CREATE TABLE queue_dispatch_claims ( + entry_id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + message_json TEXT NOT NULL, + accepted INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMP NOT NULL + ) + `) + if err != nil { + t.Fatal(err) + } + for _, legacy := range []struct { + msg QueuedMessage + accepted int + }{ + { + msg: QueuedMessage{ + ID: "legacy-dispatch-pending", SessionID: "legacy-session-pending", TaskID: "legacy-task", + Content: "legacy pending prompt", QueuedBy: QueuedByUser, + }, + }, + { + msg: QueuedMessage{ + ID: "legacy-dispatch-accepted", SessionID: "legacy-session-accepted", TaskID: "legacy-task", + Content: "legacy accepted prompt", QueuedBy: QueuedByUser, + }, + accepted: 1, + }, + } { + messageJSON, err := json.Marshal(legacy.msg) + if err != nil { + t.Fatal(err) + } + if _, err := repo.db.ExecContext(ctx, ` + INSERT INTO queue_dispatch_claims + (entry_id, session_id, message_json, accepted, created_at) + VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP) + `, legacy.msg.ID, legacy.msg.SessionID, string(messageJSON), legacy.accepted); err != nil { + t.Fatal(err) + } + } + + pending, err := repo.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 2 { + t.Fatalf("migrated dispatch count = %d, want 2", len(pending)) + } + attemptIDs := make(map[string]struct{}, len(pending)) + for i := range pending { + attemptID := pending[i].Message.dispatchAttemptID + if attemptID == "" { + t.Fatalf("migrated dispatch = %#v, want non-empty attempt id", pending[i]) + } + if _, exists := attemptIDs[attemptID]; exists { + t.Fatalf("duplicate migrated dispatch attempt id %q", attemptID) + } + attemptIDs[attemptID] = struct{}{} + if pending[i].Accepted { + err = repo.DeletePendingQueueDispatch(ctx, &pending[i].Message) + } else { + err = repo.Restore(ctx, &pending[i].Message, 0) + } + if err != nil { + t.Fatalf("settle migrated dispatch %q: %v", pending[i].Message.ID, err) + } + } + pending, err = repo.ListPendingQueueDispatches(ctx) + if err != nil || len(pending) != 0 { + t.Fatalf("pending after migrated settlement = %#v, err=%v", pending, err) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_edit_lease.go b/apps/backend/internal/orchestrator/messagequeue/repository_edit_lease.go new file mode 100644 index 0000000000..5def17b16b --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_edit_lease.go @@ -0,0 +1,159 @@ +package messagequeue + +import ( + "context" + "fmt" + "time" + + "github.com/jmoiron/sqlx" +) + +const editLeaseSchema = ` + CREATE TABLE IF NOT EXISTS queue_edit_leases ( + session_id TEXT NOT NULL, + entry_id TEXT NOT NULL, + lease_id TEXT NOT NULL, + expires_at TIMESTAMP NOT NULL, + PRIMARY KEY (session_id, entry_id) + ) +` + +type durableEditLeaseMutationRepository interface { + updateContentAndMetadataWithLease( + ctx context.Context, + sessionID, entryID, leaseID, content string, + attachments []MessageAttachment, + metadataUpdates map[string]interface{}, + queuedBy string, + ) error +} + +func (r *sqliteRepository) ensureEditLeaseSchema(ctx context.Context) error { + if _, err := r.db.ExecContext(ctx, editLeaseSchema); err != nil { + return fmt.Errorf("ensure queue edit lease schema: %w", err) + } + return nil +} + +func (r *sqliteRepository) acquireEditLease(ctx context.Context, lease *QueueEditLease) error { + tx, err := r.beginSessionMutationTx(ctx, lease.SessionID, "acquire edit lease") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + now := time.Now().UTC() + if _, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM queue_edit_leases + WHERE session_id = ? AND entry_id = ? AND expires_at <= ? + `), lease.SessionID, lease.EntryID, now); err != nil { + return fmt.Errorf("expire queue edit lease: %w", err) + } + result, err := tx.ExecContext(ctx, tx.Rebind(` + INSERT INTO queue_edit_leases (session_id, entry_id, lease_id, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(session_id, entry_id) DO NOTHING + `), lease.SessionID, lease.EntryID, lease.LeaseID, lease.ExpiresAt) + if err != nil { + return fmt.Errorf("acquire queue edit lease: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("count acquired queue edit leases: %w", err) + } + if affected != 1 { + return ErrEditConflict + } + return tx.Commit() +} + +func (r *sqliteRepository) renewEditLease(ctx context.Context, lease *QueueEditLease) error { + tx, err := r.beginSessionMutationTx(ctx, lease.SessionID, "renew edit lease") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + result, err := tx.ExecContext(ctx, tx.Rebind(` + UPDATE queue_edit_leases SET expires_at = ? + WHERE session_id = ? AND entry_id = ? AND lease_id = ? AND expires_at > ? + `), lease.ExpiresAt, lease.SessionID, lease.EntryID, lease.LeaseID, time.Now().UTC()) + if err != nil { + return fmt.Errorf("renew queue edit lease: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("count renewed queue edit leases: %w", err) + } + if affected != 1 { + return ErrEditLeaseNotFound + } + return tx.Commit() +} + +func (r *sqliteRepository) releaseEditLease(ctx context.Context, sessionID, entryID, leaseID string) error { + tx, err := r.beginSessionMutationTx(ctx, sessionID, "release edit lease") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + if _, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM queue_edit_leases WHERE session_id = ? AND entry_id = ? AND lease_id = ? + `), sessionID, entryID, leaseID); err != nil { + return fmt.Errorf("release queue edit lease: %w", err) + } + return tx.Commit() +} + +func (r *sqliteRepository) editLeaseBlocksEntryTx(ctx context.Context, tx *sqlx.Tx, sessionID, entryID string) (bool, error) { + var blocked bool + if err := tx.GetContext(ctx, &blocked, tx.Rebind(` + SELECT EXISTS ( + SELECT 1 FROM queue_edit_leases + WHERE session_id = ? AND entry_id = ? AND expires_at > ? + ) + `), sessionID, entryID, time.Now().UTC()); err != nil { + return false, fmt.Errorf("check queue edit lease: %w", err) + } + return blocked, nil +} + +func (r *sqliteRepository) deleteEditLeaseForEntryTx(ctx context.Context, tx *sqlx.Tx, sessionID, entryID string) error { + if _, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM queue_edit_leases WHERE session_id = ? AND entry_id = ? + `), sessionID, entryID); err != nil { + return fmt.Errorf("invalidate replaced queue edit lease: %w", err) + } + return nil +} + +func deleteEditLeasesForSessionIDsTx( + ctx context.Context, + tx *sqlx.Tx, + db *sqlx.DB, + sessionIDs []string, +) error { + for _, sessionID := range sessionIDs { + if _, err := tx.ExecContext(ctx, db.Rebind(` + DELETE FROM queue_edit_leases WHERE session_id = ? + `), sessionID); err != nil { + return fmt.Errorf("invalidate queue edit leases for session %s: %w", sessionID, err) + } + } + return nil +} + +func (r *sqliteRepository) deleteEditLeasesForSessionTx(ctx context.Context, tx *sqlx.Tx, sessionID string) error { + return deleteEditLeasesForSessionIDsTx(ctx, tx, r.db, []string{sessionID}) +} + +func (r *sqliteRepository) deleteEditLeasesForTransferTx( + ctx context.Context, + tx *sqlx.Tx, + oldSessionID, newSessionID string, +) error { + if _, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM queue_edit_leases WHERE session_id IN (?, ?) + `), oldSessionID, newSessionID); err != nil { + return fmt.Errorf("invalidate transferred queue edit leases: %w", err) + } + return nil +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_memory.go b/apps/backend/internal/orchestrator/messagequeue/repository_memory.go index b271373e94..47d99dcb65 100644 --- a/apps/backend/internal/orchestrator/messagequeue/repository_memory.go +++ b/apps/backend/internal/orchestrator/messagequeue/repository_memory.go @@ -18,7 +18,9 @@ type memoryRepository struct { nextPosition map[string]int64 // sessionID -> monotonic counter pendingMoves map[string]*PendingMove generation map[string]int64 + sessionGeneration map[string]int64 sendNowGeneration map[string]int64 + sendNowClaims map[string]map[string]struct{} autoRun map[string]bool autoRunIncarnation map[string]string autoMergeOverrides map[QueueSessionIdentity]AutoMergeOverride @@ -34,7 +36,9 @@ func NewMemoryRepository() Repository { nextPosition: make(map[string]int64), pendingMoves: make(map[string]*PendingMove), generation: make(map[string]int64), + sessionGeneration: make(map[string]int64), sendNowGeneration: make(map[string]int64), + sendNowClaims: make(map[string]map[string]struct{}), autoRun: make(map[string]bool), autoRunIncarnation: make(map[string]string), autoMergeOverrides: make(map[QueueSessionIdentity]AutoMergeOverride), @@ -58,6 +62,7 @@ func (r *memoryRepository) clearSessionStateLocked(sessionID string) { delete(r.nextPosition, sessionID) delete(r.pendingMoves, sessionID) delete(r.sendNowGeneration, sessionID) + delete(r.sendNowClaims, sessionID) delete(r.autoRun, sessionID) delete(r.autoRunIncarnation, sessionID) delete(r.statusGeneration, sessionID) @@ -167,23 +172,32 @@ func (r *memoryRepository) LifecycleGeneration(_ context.Context, taskID string) return r.generation[taskID], nil } +// SessionGeneration returns the current destructive-mutation generation for a +// session. +func (r *memoryRepository) SessionGeneration(_ context.Context, sessionID string) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + return r.sessionGeneration[sessionID], nil +} + // PurgeTask removes all task rows and advances its generation. func (r *memoryRepository) PurgeTask(_ context.Context, taskID string) (int, error) { r.mu.Lock() defer r.mu.Unlock() removed := 0 + affectedSessions := make(map[string]struct{}) for sessionID, list := range r.entries { kept := list[:0] for _, msg := range list { if msg.TaskID == taskID { removed++ + affectedSessions[sessionID] = struct{}{} continue } kept = append(kept, msg) } if len(kept) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) continue } r.entries[sessionID] = kept @@ -191,8 +205,12 @@ func (r *memoryRepository) PurgeTask(_ context.Context, taskID string) (int, err for sessionID, move := range r.pendingMoves { if move != nil && move.TaskID == taskID { delete(r.pendingMoves, sessionID) + affectedSessions[sessionID] = struct{}{} } } + for sessionID := range affectedSessions { + r.sessionGeneration[sessionID]++ + } r.generation[taskID]++ return removed, nil } @@ -300,6 +318,13 @@ func (r *memoryRepository) RestoreForSession( } func (r *memoryRepository) restoreLocked(msg *QueuedMessage, maxPerSession int) error { + if msg == nil { + return ErrEntryNotFound + } + if msg.reservationGenerationsCaptured && + r.sessionGeneration[msg.SessionID] != msg.reservationSessionGeneration { + return ErrQueueDispatchClaimChanged + } list := r.entries[msg.SessionID] if maxPerSession > 0 && len(list) >= maxPerSession { return ErrQueueFull @@ -310,11 +335,11 @@ func (r *memoryRepository) restoreLocked(msg *QueuedMessage, maxPerSession int) if msg.QueuedAt.IsZero() { msg.QueuedAt = time.Now().UTC() } - clone := *msg + clone := cloneQueuedMessage(msg) index := sort.Search(len(list), func(i int) bool { return list[i].Position > clone.Position }) list = append(list, nil) copy(list[index+1:], list[index:]) - list[index] = &clone + list[index] = clone r.entries[msg.SessionID] = list if clone.Position > r.nextPosition[msg.SessionID] { r.nextPosition[msg.SessionID] = clone.Position @@ -397,8 +422,8 @@ func (r *memoryRepository) insertLocked(msg *QueuedMessage, maxPerSession int) e } r.nextPosition[msg.SessionID]++ msg.Position = r.nextPosition[msg.SessionID] - clone := *msg - r.entries[msg.SessionID] = append(list, &clone) + clone := cloneQueuedMessage(msg) + r.entries[msg.SessionID] = append(list, clone) return nil } @@ -444,6 +469,13 @@ func (r *memoryRepository) requeuePreservingFIFOLocked( msg *QueuedMessage, expectedIncarnationID string, ) error { + if msg == nil { + return ErrEntryNotFound + } + if msg.reservationGenerationsCaptured && + r.sessionGeneration[msg.SessionID] != msg.reservationSessionGeneration { + return ErrQueueDispatchClaimChanged + } if msg.IsReservedLifecycleDelivery() { return r.releaseLifecycleReservationForRetryLocked(msg, expectedIncarnationID) } @@ -464,8 +496,8 @@ func (r *memoryRepository) requeuePreservingFIFOLocked( existing.Content = msg.Content existing.Model = msg.Model existing.PlanMode = msg.PlanMode - existing.Attachments = msg.Attachments - existing.Metadata = msg.Metadata + existing.Attachments = append([]MessageAttachment(nil), msg.Attachments...) + existing.Metadata = copyMessageMetadata(msg.Metadata, 0) if msg.QueuedAt.IsZero() { existing.QueuedAt = time.Now().UTC() } else { @@ -483,9 +515,9 @@ func (r *memoryRepository) requeuePreservingFIFOLocked( msg.QueuedAt = time.Now().UTC() } msg.Position = r.nextRequeuePositionLocked(msg.SessionID, list) - clone := *msg + clone := cloneQueuedMessage(msg) newList := make([]*QueuedMessage, 0, len(list)+1) - newList = append(newList, &clone) + newList = append(newList, clone) newList = append(newList, list...) r.entries[msg.SessionID] = newList return nil @@ -540,7 +572,9 @@ func (r *memoryRepository) hasPendingCoalescedSuccessorLocked( func (r *memoryRepository) nextRequeuePositionLocked(sessionID string, list []*QueuedMessage) int64 { if len(list) == 0 { - r.nextPosition[sessionID] = 1 + if r.nextPosition[sessionID] < 1 { + r.nextPosition[sessionID] = 1 + } return 1 } minPos := list[0].Position @@ -589,8 +623,8 @@ func (r *memoryRepository) appendOrInsertTailLocked(sessionID, taskID, content, tail := list[len(list)-1] if tail.QueuedBy == queuedBy { tail.Content = tail.Content + "\n\n---\n\n" + content - out := *tail - return &out, true, nil + out := cloneQueuedMessage(tail) + return out, true, nil } } @@ -607,7 +641,7 @@ func (r *memoryRepository) appendOrInsertTailLocked(sessionID, taskID, content, if err := r.insertLocked(msg, maxPerSession); err != nil { return nil, false, err } - return msg, false, nil + return cloneQueuedMessage(msg), false, nil } // InsertOrReplaceByCoalesceKey replaces an entry with the same session/queued_by/coalesce key, or inserts when allowInsert is set. @@ -661,11 +695,11 @@ func (r *memoryRepository) insertOrReplaceByCoalesceKeyLocked( existing.Content = msg.Content existing.Model = msg.Model existing.PlanMode = msg.PlanMode - existing.Attachments = msg.Attachments - existing.Metadata = msg.Metadata + existing.Attachments = append([]MessageAttachment(nil), msg.Attachments...) + existing.Metadata = copyMessageMetadata(msg.Metadata, 0) existing.QueuedAt = msg.QueuedAt - out := *existing - return &out, true, nil + out := cloneQueuedMessage(existing) + return out, true, nil } if !allowInsert { return nil, false, ErrEntryNotFound @@ -676,7 +710,7 @@ func (r *memoryRepository) insertOrReplaceByCoalesceKeyLocked( if err := r.insertLocked(msg, maxPerSession); err != nil { return nil, false, err } - return msg, false, nil + return cloneQueuedMessage(msg), false, nil } func (r *memoryRepository) InsertOrReplaceLifecycleByCoalesceKey( _ context.Context, @@ -750,7 +784,7 @@ func (r *memoryRepository) insertOrReplaceLifecycleLocked( if err := r.insertLocked(msg, maxPerSession); err != nil { return nil, false, err } - return msg, false, nil + return cloneQueuedMessage(msg), false, nil } // ListBySession returns all entries for a session ordered by position ascending. @@ -760,7 +794,7 @@ func (r *memoryRepository) ListBySession(_ context.Context, sessionID string) ([ list := r.entries[sessionID] out := make([]QueuedMessage, len(list)) for i, m := range list { - out[i] = *m + out[i] = *cloneQueuedMessage(m) } sort.Slice(out, func(i, j int) bool { return out[i].Position < out[j].Position }) return out, nil @@ -789,6 +823,19 @@ func (r *memoryRepository) ListDurableLifecycleEntries(_ context.Context) ([]Que return out, nil } +func (r *memoryRepository) FindByID(_ context.Context, entryID string) (*QueuedMessage, error) { + r.mu.Lock() + defer r.mu.Unlock() + for _, list := range r.entries { + for _, message := range list { + if message.ID == entryID { + return cloneQueuedMessage(message), nil + } + } + } + return nil, ErrEntryNotFound +} + // CountBySession returns the number of entries for a session. func (r *memoryRepository) CountBySession(_ context.Context, sessionID string) (int, error) { r.mu.Lock() @@ -796,6 +843,24 @@ func (r *memoryRepository) CountBySession(_ context.Context, sessionID string) ( return len(r.entries[sessionID]), nil } +func lowestPositionIndex(list []*QueuedMessage) int { + index := 0 + for i := 1; i < len(list); i++ { + if list[i].Position < list[index].Position { + index = i + } + } + return index +} + +func (r *memoryRepository) captureReservationLocked(msg *QueuedMessage) { + msg.reservationSessionGeneration = r.sessionGeneration[msg.SessionID] + if msg.TaskID != "" { + msg.reservationLifecycleGeneration = r.generation[msg.TaskID] + } + msg.reservationGenerationsCaptured = true +} + // TakeHead atomically returns and deletes the lowest-position entry for the session. func (r *memoryRepository) TakeHead(_ context.Context, sessionID string) (*QueuedMessage, error) { r.mu.Lock() @@ -804,14 +869,15 @@ func (r *memoryRepository) TakeHead(_ context.Context, sessionID string) (*Queue if len(list) == 0 { return nil, nil } - head := list[0] - r.entries[sessionID] = list[1:] + headIndex := lowestPositionIndex(list) + head := list[headIndex] + r.entries[sessionID] = append(list[:headIndex], list[headIndex+1:]...) if len(r.entries[sessionID]) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } - out := *head - return &out, nil + out := cloneQueuedMessage(head) + r.captureReservationLocked(out) + return out, nil } // ReserveHead returns the lowest-position entry, deleting ordinary rows and reserving durable lifecycle rows. @@ -830,27 +896,34 @@ func (r *memoryRepository) reserveHeadLocked( if len(list) == 0 { return nil } - head := list[0] - out := *head + headIndex := lowestPositionIndex(list) + head := list[headIndex] + out := cloneQueuedMessage(head) + r.captureReservationLocked(out) if head.IsDurableLifecycle() || retainOrdinary { reservationIncarnation := lifecycleReservationIncarnation(head.Metadata) - if identity != nil && reservationIncarnation != "" && - reservationIncarnation != identity.SessionIncarnationID { - r.removeEntryLocked(sessionID, 0) + if identity != nil && reservationIncarnation != "" && reservationIncarnation != identity.SessionIncarnationID { + r.removeEntryLocked(sessionID, headIndex) return r.reserveHeadLocked(sessionID, identity, retainOrdinary) } - out.Metadata = clearReservedMetadata(head.Metadata) + // Mirror the SQLite reservation: the stored row is flagged in flight so + // queue status stops listing it, while the returned copy keeps the + // unmarked metadata a requeue would write back. + out.Metadata = clearReservedMetadata(out.Metadata) out.reservedLifecycleDelivery = head.IsDurableLifecycle() if identity != nil { out.reservationIdentity = *identity head.Metadata = markReservedMetadataForIncarnation(out.Metadata, identity.SessionIncarnationID) + } else if head.IsDurableLifecycle() { + out.lifecycleReservationID = uuid.NewString() + head.Metadata = markReservedMetadata(out.Metadata, out.lifecycleReservationID) } else { head.Metadata = markReservedMetadata(out.Metadata) } - return &out + return out } - r.removeEntryLocked(sessionID, 0) - return &out + r.removeEntryLocked(sessionID, headIndex) + return out } func (r *memoryRepository) removeEntryLocked(sessionID string, index int) { @@ -858,7 +931,6 @@ func (r *memoryRepository) removeEntryLocked(sessionID string, index int) { r.entries[sessionID] = append(list[:index], list[index+1:]...) if len(r.entries[sessionID]) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } } @@ -997,20 +1069,39 @@ func (r *memoryRepository) ReserveHeadForDeliveryIfAutoRunForSession( return r.reserveHeadLocked(identity.SessionID, &identity, true), true, nil } -// AcknowledgeByID removes a reserved durable entry after executor acceptance. +// AcknowledgeReserved removes only the exact lifecycle delivery attempt. +func (r *memoryRepository) AcknowledgeReserved(_ context.Context, reserved *QueuedMessage) error { + r.mu.Lock() + defer r.mu.Unlock() + list := r.entries[reserved.SessionID] + for i, msg := range list { + if msg.ID != reserved.ID { + continue + } + storedReservationID, _ := msg.Metadata[metadataLifecycleReservationID].(string) + if reserved.lifecycleReservationID == "" || storedReservationID != reserved.lifecycleReservationID { + return ErrLifecycleReservationChanged + } + r.entries[reserved.SessionID] = append(list[:i], list[i+1:]...) + if len(r.entries[reserved.SessionID]) == 0 { + delete(r.entries, reserved.SessionID) + } + return nil + } + return ErrEntryNotFound +} + func (r *memoryRepository) AcknowledgeByID(_ context.Context, sessionID, entryID string) error { r.mu.Lock() defer r.mu.Unlock() - list := r.entries[sessionID] - for i, msg := range list { + for index, msg := range r.entries[sessionID] { if msg.ID != entryID { continue } - r.entries[sessionID] = append(list[:i], list[i+1:]...) - if len(r.entries[sessionID]) == 0 { - delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) + if !msg.IsReservedInFlight() { + return ErrEntryNotFound } + r.removeEntryLocked(sessionID, index) return nil } return ErrEntryNotFound @@ -1127,10 +1218,10 @@ func (r *memoryRepository) takeByIDLocked(sessionID, entryID string) *QueuedMess r.entries[sessionID] = append(list[:i], list[i+1:]...) if len(r.entries[sessionID]) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } - out := *message - return &out + out := cloneQueuedMessage(message) + r.captureReservationLocked(out) + return out } return nil } @@ -1182,6 +1273,7 @@ func (r *memoryRepository) claimSendNowLocked( if err := validateSendNowSnapshot(selected, expected); err != nil { return nil, ErrSendNowClaimChanged } + sessionGeneration := r.sessionGeneration[sessionID] sources := cloneSendNowSources(selected) bindSendNowLifecycleReservations(sources, identity) @@ -1195,7 +1287,6 @@ func (r *memoryRepository) claimSendNowLocked( generations[source.TaskID] = r.generation[source.TaskID] } } - remaining := make([]*QueuedMessage, 0, len(list)) for _, entry := range list { if _, ok := requested[entry.ID]; !ok { @@ -1204,7 +1295,7 @@ func (r *memoryRepository) claimSendNowLocked( } if entry.IsDurableLifecycle() { entry.Metadata = markReservedMetadataForIncarnation( - entry.Metadata, + markReservedMetadata(entry.Metadata, uuid.NewString()), identity.SessionIncarnationID, ) remaining = append(remaining, entry) @@ -1212,18 +1303,28 @@ func (r *memoryRepository) claimSendNowLocked( } if len(remaining) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } else { r.entries[sessionID] = remaining } + claimIDs := make(map[string]struct{}, len(sources)) + for _, source := range sources { + if source.IsDurableLifecycle() { + claimIDs[source.ID] = struct{}{} + } + } + if len(claimIDs) > 0 { + r.sendNowClaims[sessionID] = claimIDs + } r.setAutoRunLocked(identityPointer(identity), sessionID, true) r.sendNowGeneration[sessionID]++ return &SendNowClaim{ + ClaimID: uuid.NewString(), Identity: identity, OperationGeneration: r.sendNowGeneration[sessionID], Sources: sources, Dispatch: *envelope, SourceGenerations: generations, + SessionGeneration: sessionGeneration, }, nil } @@ -1266,6 +1367,16 @@ func (r *memoryRepository) RestoreSendNowClaim(ctx context.Context, claim *SendN if err != nil { return err } + generations := make(map[string]int64) + for _, source := range claim.Sources { + if source.TaskID != "" { + generations[source.TaskID] = r.generation[source.TaskID] + } + } + sessionChanged := claim.SessionGeneration != r.sessionGeneration[sessionID] + if sessionChanged && !sendNowClaimSourcesAllInvalidated(claim, generations) { + return ErrSendNowClaimChanged + } list := r.entries[sessionID] existing := make(map[string]*QueuedMessage, len(list)) for _, entry := range list { @@ -1294,6 +1405,7 @@ func (r *memoryRepository) RestoreSendNowClaim(ctx context.Context, claim *SendN } sort.Slice(list, func(i, j int) bool { return list[i].Position < list[j].Position }) r.entries[sessionID] = list + delete(r.sendNowClaims, sessionID) return nil } @@ -1339,12 +1451,25 @@ func (r *memoryRepository) AcknowledgeSendNowClaim(ctx context.Context, claim *S if err != nil { return err } + generations := make(map[string]int64) + for _, source := range claim.Sources { + if source.TaskID != "" { + generations[source.TaskID] = r.generation[source.TaskID] + } + } + if claim.SessionGeneration != r.sessionGeneration[sessionID] && + !sendNowClaimSourcesAllInvalidated(claim, generations) { + return ErrSendNowClaimChanged + } requested := make(map[string]struct{}) for _, source := range claim.Sources { if source.SessionID != sessionID { return ErrSendNowClaimChanged } - if !source.IsDurableLifecycle() { + if source.TaskID != "" { + generations[source.TaskID] = r.generation[source.TaskID] + } + if !source.IsDurableLifecycle() || sendNowSourceGenerationChanged(claim, source, generations[source.TaskID]) { continue } requested[source.ID] = struct{}{} @@ -1366,10 +1491,10 @@ func (r *memoryRepository) AcknowledgeSendNowClaim(ctx context.Context, claim *S } if len(remaining) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } else { r.entries[sessionID] = remaining } + delete(r.sendNowClaims, sessionID) return nil } @@ -1413,10 +1538,25 @@ func sameQueuedMessageContent(left, right *QueuedMessage) bool { } leftCopy := cloneQueuedMessage(left) rightCopy := cloneQueuedMessage(right) + // PostgreSQL stores TIMESTAMP values at microsecond precision. Normalize + // both snapshots before comparing so a caller's nanosecond timestamp does + // not look like a queue mutation after a round trip through the database. + leftCopy.QueuedAt = leftCopy.QueuedAt.Truncate(time.Microsecond) + rightCopy.QueuedAt = rightCopy.QueuedAt.Truncate(time.Microsecond) leftCopy.Metadata = clearReservedMetadata(leftCopy.Metadata) rightCopy.Metadata = clearReservedMetadata(rightCopy.Metadata) leftCopy.reservedLifecycleDelivery = false rightCopy.reservedLifecycleDelivery = false + leftCopy.dispatchAttemptID = "" + rightCopy.dispatchAttemptID = "" + leftCopy.lifecycleReservationID = "" + rightCopy.lifecycleReservationID = "" + leftCopy.reservationSessionGeneration = 0 + rightCopy.reservationSessionGeneration = 0 + leftCopy.reservationLifecycleGeneration = 0 + rightCopy.reservationLifecycleGeneration = 0 + leftCopy.reservationGenerationsCaptured = false + rightCopy.reservationGenerationsCaptured = false return reflect.DeepEqual(leftCopy, rightCopy) } @@ -1466,7 +1606,7 @@ func (r *memoryRepository) updateContentAndMetadataLocked(sessionID, entryID, co return ErrEntryNotFound } m.Content = content - m.Attachments = attachments + m.Attachments = append([]MessageAttachment(nil), attachments...) m.Metadata = applyMetadataUpdates(m.Metadata, metadataUpdates) return nil } @@ -1544,8 +1684,7 @@ func (r *memoryRepository) mergeIntoAboveLocked(sessionID, sourceID, queuedBy st if len(r.entries[sessionID]) == 0 { delete(r.nextPosition, sessionID) } - merged := *target - return &merged, nil + return cloneQueuedMessage(target), nil } // AutoMergeIntoAbove folds one exact source into its immediate compatible @@ -1812,7 +1951,6 @@ func (r *memoryRepository) deleteByIDLocked(sessionID, entryID string) error { r.entries[sessionID] = append(list[:i], list[i+1:]...) if len(r.entries[sessionID]) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } return nil } @@ -1870,10 +2008,23 @@ func (r *memoryRepository) deleteAllBySessionLocked(sessionID string) (int, erro } if len(kept) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } else { r.entries[sessionID] = kept } + r.sessionGeneration[sessionID]++ + return removed, nil +} + +// PurgeSession removes all queue rows for a deleted session, including +// reserved lifecycle deliveries, and clears its pending move and policy. +func (r *memoryRepository) PurgeSession(_ context.Context, sessionID string) (int, error) { + r.mu.Lock() + defer r.mu.Unlock() + removed := len(r.entries[sessionID]) + delete(r.entries, sessionID) + delete(r.pendingMoves, sessionID) + delete(r.autoRun, sessionID) + r.sessionGeneration[sessionID]++ return removed, nil } @@ -1919,8 +2070,10 @@ func (r *memoryRepository) transferSessionLocked( if oldSessionID == newSessionID { return nil } + r.sessionGeneration[oldSessionID]++ + r.sessionGeneration[newSessionID]++ for _, entry := range r.entries[oldSessionID] { - if entry.IsReservedInFlight() { + if entry.IsReservedInFlight() && !r.sendNowClaimContainsLocked(oldSessionID, entry.ID) { return ErrQueueChanged } } @@ -1946,18 +2099,20 @@ func (r *memoryRepository) transferSessionLocked( m.Position += destMax } r.entries[newSessionID] = append(r.entries[newSessionID], list...) - // Recompute nextPosition for the destination so future inserts keep - // monotonic ordering. - var maxPos int64 + // Keep the destination high-water mark even when its physical queue + // was drained before this transfer. for _, m := range r.entries[newSessionID] { - if m.Position > maxPos { - maxPos = m.Position + if m.Position > r.nextPosition[newSessionID] { + r.nextPosition[newSessionID] = m.Position } } - r.nextPosition[newSessionID] = maxPos delete(r.entries, oldSessionID) - delete(r.nextPosition, oldSessionID) } + if claimIDs := r.sendNowClaims[oldSessionID]; len(claimIDs) > 0 { + r.sendNowClaims[newSessionID] = claimIDs + delete(r.sendNowClaims, oldSessionID) + } + delete(r.nextPosition, oldSessionID) if move, ok := r.pendingMoves[oldSessionID]; ok { if destination, exists := r.identities[newSessionID]; exists { move.SessionIncarnationID = destination.SessionIncarnationID @@ -1973,10 +2128,16 @@ func (r *memoryRepository) transferSessionLocked( return nil } +func (r *memoryRepository) sendNowClaimContainsLocked(sessionID, entryID string) bool { + _, ok := r.sendNowClaims[sessionID][entryID] + return ok +} + // ReplaceSession replaces a session's queue with the supplied snapshot. func (r *memoryRepository) ReplaceSession(_ context.Context, sessionID string, entries []QueuedMessage, pendingMove *PendingMove) error { r.mu.Lock() defer r.mu.Unlock() + r.sessionGeneration[sessionID]++ return r.replaceSessionLocked(sessionID, entries, pendingMove) } @@ -1995,20 +2156,24 @@ func (r *memoryRepository) ReplaceSessionForIdentity(_ context.Context, identity func (r *memoryRepository) replaceSessionLocked(sessionID string, entries []QueuedMessage, pendingMove *PendingMove) error { if len(entries) == 0 { delete(r.entries, sessionID) - delete(r.nextPosition, sessionID) } else { replaced := make([]*QueuedMessage, 0, len(entries)) var maxPos int64 for _, entry := range entries { - clone := entry + clone := cloneQueuedMessage(&entry) clone.SessionID = sessionID if clone.Position > maxPos { maxPos = clone.Position } - replaced = append(replaced, &clone) + replaced = append(replaced, clone) } + sort.Slice(replaced, func(i, j int) bool { + return replaced[i].Position < replaced[j].Position + }) r.entries[sessionID] = replaced - r.nextPosition[sessionID] = maxPos + if maxPos > r.nextPosition[sessionID] { + r.nextPosition[sessionID] = maxPos + } } if pendingMove == nil { delete(r.pendingMoves, sessionID) diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_memory_regression_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_memory_regression_test.go new file mode 100644 index 0000000000..d167f06e8e --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_memory_regression_test.go @@ -0,0 +1,279 @@ +package messagequeue + +import ( + "context" + "testing" + "time" +) + +func TestMemoryReplaceSessionOwnsQueueSnapshotData(t *testing.T) { + repo := NewMemoryRepository() + ctx := context.Background() + original := QueuedMessage{ + ID: "entry-1", + SessionID: "source-session", + TaskID: "task-1", + Position: 1, + Content: "original", + QueuedAt: time.Unix(1, 0).UTC(), + QueuedBy: QueuedByUser, + Attachments: []MessageAttachment{{ + AttachmentID: "attachment-1", + Name: "original.txt", + }}, + Metadata: map[string]interface{}{"source": "original"}, + } + move := &PendingMove{MoveID: "move-1", TaskID: "task-1", QueuedAt: time.Unix(2, 0).UTC()} + + if err := repo.ReplaceSession(ctx, "restored-session", []QueuedMessage{original}, move); err != nil { + t.Fatalf("replace session: %v", err) + } + original.Attachments[0].Name = "mutated.txt" + original.Metadata["source"] = "mutated" + move.MoveID = "mutated-move" + + entries, err := repo.ListBySession(ctx, "restored-session") + if err != nil { + t.Fatalf("list restored session: %v", err) + } + if len(entries) != 1 { + t.Fatalf("restored entry count = %d, want 1", len(entries)) + } + if got := entries[0].Attachments[0].Name; got != "original.txt" { + t.Fatalf("restored attachment name = %q, want original.txt", got) + } + if got := entries[0].Metadata["source"]; got != "original" { + t.Fatalf("restored metadata source = %q, want original", got) + } + storedMove, err := repo.GetPendingMove(ctx, "restored-session") + if err != nil { + t.Fatalf("get restored move: %v", err) + } + if storedMove == nil || storedMove.MoveID != "move-1" { + t.Fatalf("restored move = %#v, want move-1", storedMove) + } +} + +func TestMemoryReserveHeadUsesFIFOPositionAfterReplacement(t *testing.T) { + svc := setupService(t) + repo := svc.repo + ctx := context.Background() + if err := repo.ReplaceSession(ctx, "session-1", []QueuedMessage{ + {ID: "late", SessionID: "session-1", TaskID: "task-1", Position: 2, Content: "late", QueuedBy: QueuedByUser}, + {ID: "head", SessionID: "session-1", TaskID: "task-1", Position: 1, Content: "head", QueuedBy: QueuedByUser}, + }, nil); err != nil { + t.Fatalf("replace session: %v", err) + } + + reserved, ok := svc.ReserveQueued(ctx, "session-1") + if !ok || reserved == nil { + t.Fatalf("reserve head = %#v, %v", reserved, ok) + } + if reserved.ID != "head" { + t.Fatalf("reserved entry = %q, want head", reserved.ID) + } +} + +func TestMemoryInsertAndListOwnQueueSnapshotData(t *testing.T) { + repo := NewMemoryRepository() + ctx := context.Background() + original := &QueuedMessage{ + ID: "entry-1", + SessionID: "session-1", + TaskID: "task-1", + Content: "original", + QueuedBy: QueuedByUser, + Attachments: []MessageAttachment{{ + AttachmentID: "attachment-1", + Name: "original.txt", + }}, + Metadata: map[string]interface{}{ + "nested": map[string]interface{}{"value": "original"}, + }, + } + if err := repo.Insert(ctx, original, 0); err != nil { + t.Fatalf("insert: %v", err) + } + + original.Attachments[0].Name = "caller-mutated.txt" + original.Metadata["nested"].(map[string]interface{})["value"] = "caller-mutated" + entries, err := repo.ListBySession(ctx, original.SessionID) + if err != nil { + t.Fatalf("list after caller mutation: %v", err) + } + if got := entries[0].Attachments[0].Name; got != "original.txt" { + t.Fatalf("stored attachment after caller mutation = %q, want original.txt", got) + } + if got := entries[0].Metadata["nested"].(map[string]interface{})["value"]; got != "original" { + t.Fatalf("stored metadata after caller mutation = %q, want original", got) + } + + entries[0].Attachments[0].Name = "result-mutated.txt" + entries[0].Metadata["nested"].(map[string]interface{})["value"] = "result-mutated" + fresh, err := repo.ListBySession(ctx, original.SessionID) + if err != nil { + t.Fatalf("list after result mutation: %v", err) + } + if got := fresh[0].Attachments[0].Name; got != "original.txt" { + t.Fatalf("stored attachment after result mutation = %q, want original.txt", got) + } + if got := fresh[0].Metadata["nested"].(map[string]interface{})["value"]; got != "original" { + t.Fatalf("stored metadata after result mutation = %q, want original", got) + } + lifecycle := &QueuedMessage{ + ID: "lifecycle-1", + SessionID: "lifecycle-session", + TaskID: "task-1", + Content: "lifecycle", + QueuedBy: QueuedByWorkflow, + Attachments: []MessageAttachment{{ + AttachmentID: "lifecycle-attachment", + Name: "lifecycle.txt", + }}, + Metadata: map[string]interface{}{ + MetadataLifecycleDurable: true, + "nested": map[string]interface{}{"value": "original"}, + }, + } + if err := repo.Insert(ctx, lifecycle, 0); err != nil { + t.Fatalf("insert lifecycle: %v", err) + } + reserved, err := repo.ReserveHead(ctx, lifecycle.SessionID) + if err != nil { + t.Fatalf("reserve lifecycle: %v", err) + } + reserved.Attachments[0].Name = "result-mutated.txt" + reserved.Metadata["nested"].(map[string]interface{})["value"] = "result-mutated" + fresh, err = repo.ListBySession(ctx, lifecycle.SessionID) + if err != nil { + t.Fatalf("list after reserve result mutation: %v", err) + } + if got := fresh[0].Attachments[0].Name; got != "lifecycle.txt" { + t.Fatalf("reserved attachment after result mutation = %q, want lifecycle.txt", got) + } + if got := fresh[0].Metadata["nested"].(map[string]interface{})["value"]; got != "original" { + t.Fatalf("reserved metadata after result mutation = %q, want original", got) + } +} + +func TestMemoryAppendInsertReturnsOwnedSnapshot(t *testing.T) { + repo := NewMemoryRepository() + ctx := context.Background() + attachments := []MessageAttachment{{AttachmentID: "attachment-1", Name: "original.txt"}} + metadata := map[string]interface{}{ + "nested": map[string]interface{}{"value": "original"}, + } + + out, appended, err := repo.AppendOrInsertTail( + ctx, "session-append-snapshot", "task-1", "content", "", QueuedByUser, + false, attachments, metadata, 0, + ) + if err != nil { + t.Fatalf("append or insert: %v", err) + } + if appended { + t.Fatal("initial append or insert unexpectedly appended") + } + out.Attachments[0].Name = "mutated.txt" + out.Metadata["nested"].(map[string]interface{})["value"] = "mutated" + if got := attachments[0].Name; got != "original.txt" { + t.Fatalf("input attachment after returned snapshot mutation = %q, want original.txt", got) + } + if got := metadata["nested"].(map[string]interface{})["value"]; got != "original" { + t.Fatalf("input metadata after returned snapshot mutation = %q, want original", got) + } + + entries, err := repo.ListBySession(ctx, "session-append-snapshot") + if err != nil { + t.Fatalf("list after returned snapshot mutation: %v", err) + } + if got := entries[0].Attachments[0].Name; got != "original.txt" { + t.Fatalf("stored attachment after returned snapshot mutation = %q, want original.txt", got) + } + if got := entries[0].Metadata["nested"].(map[string]interface{})["value"]; got != "original" { + t.Fatalf("stored metadata after returned snapshot mutation = %q, want original", got) + } +} + +func TestMemoryCoalesceReplacementPreservesTimestampAndPlanMode(t *testing.T) { + repo := NewMemoryRepository() + ctx := context.Background() + coalesceKey := "coalesce-key" + original := &QueuedMessage{ + SessionID: "session-coalesce-snapshot", + TaskID: "task-1", + Content: "original", + PlanMode: false, + QueuedBy: QueuedByWorkflow, + QueuedAt: time.Unix(10, 0).UTC(), + Metadata: map[string]interface{}{MetadataCoalesceKey: coalesceKey}, + } + if err := repo.Insert(ctx, original, 0); err != nil { + t.Fatalf("insert original: %v", err) + } + replacement := &QueuedMessage{ + SessionID: "session-coalesce-snapshot", + TaskID: "task-2", + Content: "replacement", + PlanMode: true, + QueuedBy: QueuedByWorkflow, + QueuedAt: time.Unix(20, 0).UTC(), + Metadata: map[string]interface{}{MetadataCoalesceKey: coalesceKey}, + } + updated, replaced, err := repo.InsertOrReplaceByCoalesceKey( + ctx, replacement, coalesceKey, 0, true, + ) + if err != nil { + t.Fatalf("replace ordinary coalesced entry: %v", err) + } + if !replaced { + t.Fatal("ordinary coalesce unexpectedly inserted") + } + if !updated.PlanMode || !updated.QueuedAt.Equal(replacement.QueuedAt) { + t.Fatalf("ordinary replacement = plan_mode:%t queued_at:%s, want plan_mode:true queued_at:%s", + updated.PlanMode, updated.QueuedAt, replacement.QueuedAt) + } + + lifecycle := &QueuedMessage{ + SessionID: "session-lifecycle-coalesce-snapshot", + TaskID: "task-1", + Content: "original lifecycle", + PlanMode: false, + QueuedBy: QueuedByWorkflow, + QueuedAt: time.Unix(30, 0).UTC(), + Metadata: map[string]interface{}{ + MetadataCoalesceKey: coalesceKey, + MetadataLifecycleDurable: true, + MetadataLifecycleGeneration: int64(0), + }, + } + if err := repo.Insert(ctx, lifecycle, 0); err != nil { + t.Fatalf("insert lifecycle original: %v", err) + } + lifecycleReplacement := &QueuedMessage{ + SessionID: "session-lifecycle-coalesce-snapshot", + TaskID: "task-2", + Content: "replacement lifecycle", + PlanMode: true, + QueuedBy: QueuedByWorkflow, + QueuedAt: time.Unix(40, 0).UTC(), + Metadata: map[string]interface{}{ + MetadataCoalesceKey: coalesceKey, + MetadataLifecycleDurable: true, + MetadataLifecycleGeneration: int64(0), + }, + } + updated, replaced, err = repo.InsertOrReplaceLifecycleByCoalesceKey( + ctx, lifecycleReplacement, coalesceKey, 0, true, + ) + if err != nil { + t.Fatalf("replace lifecycle coalesced entry: %v", err) + } + if !replaced { + t.Fatal("lifecycle coalesce unexpectedly inserted") + } + if !updated.PlanMode || !updated.QueuedAt.Equal(lifecycleReplacement.QueuedAt) { + t.Fatalf("lifecycle replacement = plan_mode:%t queued_at:%s, want plan_mode:true queued_at:%s", + updated.PlanMode, updated.QueuedAt, lifecycleReplacement.QueuedAt) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_postgres_durability_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_postgres_durability_test.go new file mode 100644 index 0000000000..872ad522f7 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_postgres_durability_test.go @@ -0,0 +1,478 @@ +package messagequeue + +import ( + "context" + "errors" + "testing" +) + +func TestPostgresRepository_DurableQueueRecoveryTables(t *testing.T) { + ctx := context.Background() + repo := newTestPostgresRepo(t) + persistent := repo.(*sqliteRepository) + + t.Run("attachment cleanup replays schema and upsert", func(t *testing.T) { + cleanup := AttachmentCleanup{ + SessionID: "cleanup-session", EntryID: "cleanup-entry", OperationID: "cleanup-op", + TaskID: "cleanup-task", OwnerID: "owner-1", LeaseID: "lease-1", + Attachments: []MessageAttachment{{Name: "first.txt", Data: "first"}}, + } + if err := persistent.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + cleanup.OwnerID = "owner-2" + cleanup.Attachments = []MessageAttachment{{Name: "second.txt", Data: "second"}} + if err := persistent.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + for range 2 { + cleanups, err := persistent.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(cleanups) != 1 || cleanups[0].OwnerID != "owner-2" || cleanups[0].Attachments[0].Name != "second.txt" { + t.Fatalf("attachment cleanups = %#v", cleanups) + } + } + if err := persistent.DeleteAttachmentCleanup(ctx, cleanup.SessionID, cleanup.EntryID, cleanup.OperationID); err != nil { + t.Fatal(err) + } + }) + + t.Run("session transfer compensation replays schema and upsert", func(t *testing.T) { + compensation := SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: "transfer-task", FromSessionID: "transfer-old", ToSessionID: "transfer-new", + EntryIDs: []string{"entry-1"}, + } + if err := persistent.UpsertSessionTransferCompensation(ctx, compensation); err != nil { + t.Fatal(err) + } + compensation.EntryIDs = []string{"entry-1", "entry-2"} + if err := persistent.UpsertSessionTransferCompensation(ctx, compensation); err != nil { + t.Fatal(err) + } + compensations, err := persistent.ListSessionTransferCompensations(ctx) + if err != nil { + t.Fatal(err) + } + if len(compensations) != 1 || len(compensations[0].EntryIDs) != 2 || compensations[0].EntryIDs[1] != "entry-2" { + t.Fatalf("session transfer compensations = %#v", compensations) + } + if err := persistent.DeleteSessionTransferCompensation(ctx, compensation.OperationID, compensation.TaskID, compensation.FromSessionID, compensation.ToSessionID); err != nil { + t.Fatal(err) + } + }) + + t.Run("send now claim records acceptance", func(t *testing.T) { + source := insertTestEntry(t, repo, "send-now-session", "send-now-task", "prompt", QueuedByUser, nil, nil) + claim, err := repo.ClaimSendNow(ctx, source.SessionID, []QueuedMessage{*source}) + if err != nil { + t.Fatal(err) + } + if err := persistent.MarkPendingSendNowClaimAccepted(ctx, claim); err != nil { + t.Fatal(err) + } + claims, err := persistent.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(claims) != 1 || !claims[0].Accepted || claims[0].Claim.Dispatch.ID != claim.Dispatch.ID { + t.Fatalf("pending Send Now claims = %#v", claims) + } + if err := repo.AcknowledgeSendNowClaim(ctx, claim); err != nil { + t.Fatal(err) + } + }) + + t.Run("ordinary dispatch claim records acceptance", func(t *testing.T) { + source := insertTestEntry(t, repo, "dispatch-session", "dispatch-task", "prompt", QueuedByUser, nil, nil) + reserved, enabled, err := repo.ReserveHeadIfAutoRun(ctx, source.SessionID) + if err != nil { + t.Fatal(err) + } + if !enabled || reserved == nil || reserved.ID != source.ID { + t.Fatalf("ordinary reservation = %#v, enabled=%t", reserved, enabled) + } + if err := persistent.MarkPendingQueueDispatchAccepted(ctx, reserved); err != nil { + t.Fatal(err) + } + dispatches, err := persistent.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(dispatches) != 1 || !dispatches[0].Accepted || dispatches[0].Message.ID != source.ID { + t.Fatalf("pending ordinary dispatches = %#v", dispatches) + } + if err := persistent.DeletePendingQueueDispatch(ctx, reserved); err != nil { + t.Fatal(err) + } + }) + + t.Run("task purge preserves attachment recovery obligation", func(t *testing.T) { + source := insertTestEntry(t, repo, "purge-session", "purge-task", "prompt", QueuedByUser, nil, nil) + if err := persistent.UpsertAttachmentCleanup(ctx, AttachmentCleanup{ + SessionID: source.SessionID, EntryID: source.ID, OperationID: "purge-operation", + TaskID: source.TaskID, Attachments: []MessageAttachment{{AttachmentID: "purge-attachment"}}, + }); err != nil { + t.Fatal(err) + } + if reserved, _, err := repo.ReserveHeadIfAutoRun(ctx, source.SessionID); err != nil || reserved == nil { + t.Fatalf("reserve purge source = %#v, err=%v", reserved, err) + } + if _, err := repo.PurgeTask(ctx, source.TaskID); err != nil { + t.Fatal(err) + } + dispatches, err := persistent.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + cleanups, err := persistent.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(dispatches) != 0 || len(cleanups) != 1 { + t.Fatalf("recovery rows after purge: dispatches=%#v cleanups=%#v", dispatches, cleanups) + } + }) + + t.Run("session mutation reconciles ordinary dispatch claims", func(t *testing.T) { + source := insertTestEntry(t, repo, "mutation-old", "mutation-task", "prompt", QueuedByUser, nil, nil) + if reserved, _, err := repo.ReserveHeadIfAutoRun(ctx, source.SessionID); err != nil || reserved == nil { + t.Fatalf("reserve mutation source = %#v, err=%v", reserved, err) + } + if err := repo.TransferSession(ctx, source.SessionID, "mutation-new"); err != nil { + t.Fatal(err) + } + dispatches, err := persistent.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(dispatches) != 1 || dispatches[0].Message.SessionID != "mutation-new" { + t.Fatalf("transferred dispatch claims = %#v", dispatches) + } + if err := repo.ReplaceSession(ctx, "mutation-new", nil, nil); err != nil { + t.Fatal(err) + } + dispatches, err = persistent.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(dispatches) != 0 { + t.Fatalf("dispatch claims after replacement = %#v", dispatches) + } + }) +} + +func TestPostgresAttachmentCleanupMigratesAndTransfersCurrentSession(t *testing.T) { + ctx := context.Background() + repo := newTestPostgresRepo(t).(*sqliteRepository) + if _, err := repo.db.ExecContext(ctx, `DROP TABLE IF EXISTS queue_attachment_cleanups`); err != nil { + t.Fatal(err) + } + if _, err := repo.db.ExecContext(ctx, ` + CREATE TABLE queue_attachment_cleanups ( + session_id TEXT NOT NULL, + entry_id TEXT NOT NULL, + operation_id TEXT NOT NULL DEFAULT '', + task_id TEXT NOT NULL, + owner_id TEXT NOT NULL DEFAULT '', + lease_id TEXT NOT NULL DEFAULT '', + remove_entry INTEGER NOT NULL DEFAULT 0, + claim_pending INTEGER NOT NULL DEFAULT 0, + entry_fingerprint TEXT NOT NULL DEFAULT '', + attachments_json TEXT NOT NULL DEFAULT '[]', + created_at TIMESTAMP NOT NULL, + PRIMARY KEY (session_id, entry_id, operation_id) + ) + `); err != nil { + t.Fatal(err) + } + const ( + sourceSessionID = "postgres-cleanup-source" + middleSessionID = "postgres-cleanup-middle" + finalSessionID = "postgres-cleanup-final" + ) + if _, err := repo.db.ExecContext(ctx, repo.db.Rebind(` + INSERT INTO queue_attachment_cleanups + (session_id, entry_id, operation_id, task_id, attachments_json, created_at) + VALUES (?, 'postgres-cleanup-entry', 'postgres-cleanup-operation', + 'postgres-cleanup-task', '[]', CURRENT_TIMESTAMP) + `), sourceSessionID); err != nil { + t.Fatal(err) + } + + cleanups, err := repo.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(cleanups) != 1 || cleanups[0].CurrentSessionID != sourceSessionID { + t.Fatalf("migrated PostgreSQL cleanup = %#v", cleanups) + } + if err := repo.TransferSession(ctx, sourceSessionID, middleSessionID); err != nil { + t.Fatal(err) + } + if err := repo.TransferSession(ctx, middleSessionID, finalSessionID); err != nil { + t.Fatal(err) + } + cleanups, err = repo.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(cleanups) != 1 || + cleanups[0].SessionID != sourceSessionID || + cleanups[0].CurrentSessionID != finalSessionID { + t.Fatalf("multi-hop PostgreSQL cleanup = %#v", cleanups) + } +} + +type replaceBeforeReservationRepository struct { + Repository + replacement Repository +} + +func (r *replaceBeforeReservationRepository) ReserveHeadIfAutoRun( + ctx context.Context, + sessionID string, +) (*QueuedMessage, bool, error) { + entries, err := r.replacement.ListBySession(ctx, sessionID) + if err != nil { + return nil, true, err + } + if err := r.replacement.ReplaceSession(ctx, sessionID, entries, nil); err != nil { + return nil, true, err + } + return r.Repository.ReserveHeadIfAutoRun(ctx, sessionID) +} + +func TestPostgresRepository_ReservationCapturesGenerationAfterConcurrentReplacement(t *testing.T) { + repoA, repoB, _ := newTestPostgresRepoPair(t) + ctx := context.Background() + source := insertTestEntry(t, repoA, "reservation-race-session", "reservation-race-task", "prompt", QueuedByUser, nil, nil) + service := setupService(t) + service.repo = &replaceBeforeReservationRepository{Repository: repoA, replacement: repoB} + + reserved, exists, autoRun := service.ReserveQueuedWithAutoRun(ctx, source.SessionID) + if !exists || !autoRun || reserved == nil { + t.Fatalf("reservation = %#v, exists=%t, autoRun=%t", reserved, exists, autoRun) + } + generation, err := repoA.SessionGeneration(ctx, source.SessionID) + if err != nil { + t.Fatal(err) + } + if !reserved.reservationGenerationsCaptured || reserved.reservationSessionGeneration != generation { + t.Fatalf( + "reservation generation = %d (captured=%t), want current %d", + reserved.reservationSessionGeneration, reserved.reservationGenerationsCaptured, generation, + ) + } +} + +func TestPostgresRepository_TransferFenceOwnershipIsOperationScoped(t *testing.T) { + repoA, repoB, _ := newTestPostgresRepoPair(t) + ctx := context.Background() + persistenceA := repoA.(sessionTransferCompensationRepository) + persistenceB := repoB.(sessionTransferCompensationRepository) + first := SessionTransferCompensation{ + OperationID: "transfer-first", + TaskID: "transfer-task", FromSessionID: "transfer-old", ToSessionID: "transfer-new", + EntryIDs: []string{"entry-first"}, + } + if err := persistenceA.UpsertSessionTransferCompensation(ctx, first); err != nil { + t.Fatal(err) + } + replacement := first + replacement.OperationID = "transfer-replacement" + replacement.EntryIDs = []string{"entry-replacement"} + if err := persistenceB.UpsertSessionTransferCompensation(ctx, replacement); !errors.Is(err, ErrSessionTransferInProgress) { + t.Fatalf("replacement error = %v, want %v", err, ErrSessionTransferInProgress) + } + if err := persistenceB.DeleteSessionTransferCompensation( + ctx, replacement.OperationID, first.TaskID, first.FromSessionID, first.ToSessionID, + ); !errors.Is(err, ErrSessionTransferOwnershipLost) { + t.Fatalf("stale delete error = %v, want %v", err, ErrSessionTransferOwnershipLost) + } + stored, err := persistenceA.ListSessionTransferCompensations(ctx) + if err != nil { + t.Fatal(err) + } + if len(stored) != 1 || stored[0].OperationID != first.OperationID { + t.Fatalf("session transfer compensations = %#v", stored) + } +} + +func TestPostgresRepository_ActiveSessionTransferFencesMutationsAcrossInstances(t *testing.T) { + tests := []struct { + name string + mutate func(context.Context, Repository, *QueuedMessage, *QueuedMessage) error + }{ + { + name: "edit", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + return repo.UpdateContent(ctx, first.SessionID, first.ID, "edited", nil, QueuedByUser) + }, + }, + { + name: "purge", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + _, err := repo.PurgeSession(ctx, first.SessionID) + return err + }, + }, + { + name: "purge task", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + _, err := repo.PurgeTask(ctx, first.TaskID) + return err + }, + }, + { + name: "merge", + mutate: func(ctx context.Context, repo Repository, _, second *QueuedMessage) error { + _, err := repo.MergeIntoAbove(ctx, second.SessionID, second.ID, QueuedByUser) + return err + }, + }, + { + name: "reorder", + mutate: func(ctx context.Context, repo Repository, first, second *QueuedMessage) error { + return repo.ReorderEntries(ctx, first.SessionID, []string{second.ID, first.ID}) + }, + }, + { + name: "take pending move", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + _, err := repo.TakePendingMove(ctx, first.SessionID) + return err + }, + }, + { + name: "transfer", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + return repo.TransferSession(ctx, first.SessionID, "session-third") + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + repoA, repoB, _ := newTestPostgresRepoPair(t) + ctx := context.Background() + first := insertTestEntry(t, repoA, "session-old", "task", "first", QueuedByUser, nil, nil) + second := insertTestEntry(t, repoA, "session-old", "task", "second", QueuedByUser, nil, nil) + if err := repoA.SetPendingMove(ctx, first.SessionID, &PendingMove{ + MoveID: "move", TaskID: "task", WorkflowID: "workflow", WorkflowStepID: "step", + }); err != nil { + t.Fatal(err) + } + persistence := repoA.(sessionTransferCompensationRepository) + if err := persistence.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: "task", FromSessionID: first.SessionID, ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + + if err := test.mutate(ctx, repoB, first, second); !errors.Is(err, ErrSessionTransferInProgress) { + t.Fatalf("mutation error = %v, want %v", err, ErrSessionTransferInProgress) + } + }) + } +} + +func TestPostgresRepository_ActiveSessionTransferFencesDispatchSettlementAcrossInstances(t *testing.T) { + tests := []struct { + name string + settle func(context.Context, *sqliteRepository, *QueuedMessage) error + }{ + {name: "mark accepted", settle: func(ctx context.Context, repo *sqliteRepository, msg *QueuedMessage) error { + return repo.MarkPendingQueueDispatchAccepted(ctx, msg) + }}, + {name: "delete", settle: func(ctx context.Context, repo *sqliteRepository, msg *QueuedMessage) error { + return repo.DeletePendingQueueDispatch(ctx, msg) + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + repoA, repoB, _ := newTestPostgresRepoPair(t) + ctx := context.Background() + entry := insertTestEntry(t, repoA, "session-old", "task", "first", QueuedByUser, nil, nil) + reserved, err := repoA.ReserveHead(ctx, entry.SessionID) + if err != nil { + t.Fatal(err) + } + persistentA := repoA.(sessionTransferCompensationRepository) + if err := persistentA.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + + err = test.settle(ctx, repoB.(*sqliteRepository), reserved) + if !errors.Is(err, ErrSessionTransferInProgress) { + t.Fatalf("settlement error = %v, want %v", err, ErrSessionTransferInProgress) + } + }) + } +} + +func TestPostgresRepository_ActiveSessionTransferFencesSendNowSettlementAcrossInstances(t *testing.T) { + tests := []struct { + name string + settle func(context.Context, *sqliteRepository, *SendNowClaim) error + }{ + {name: "mark accepted", settle: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.MarkPendingSendNowClaimAccepted(ctx, claim) + }}, + {name: "delete", settle: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.DeletePendingSendNowClaim(ctx, claim) + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + repoA, repoB, _ := newTestPostgresRepoPair(t) + ctx := context.Background() + entry := insertTestEntry(t, repoA, "session-old", "task", "first", QueuedByUser, nil, nil) + claim, err := repoA.ClaimSendNow(ctx, entry.SessionID, []QueuedMessage{*entry}) + if err != nil { + t.Fatal(err) + } + persistentA := repoA.(sessionTransferCompensationRepository) + if err := persistentA.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + + err = test.settle(ctx, repoB.(*sqliteRepository), claim) + if !errors.Is(err, ErrSessionTransferInProgress) { + t.Fatalf("settlement error = %v, want %v", err, ErrSessionTransferInProgress) + } + }) + } +} + +func TestPostgresRepository_ActiveSessionTransferFencesEditLeaseAcrossInstances(t *testing.T) { + repoA, repoB, _ := newTestPostgresRepoPair(t) + ctx := context.Background() + entry := insertTestEntry(t, repoA, "session-old", "task", "first", QueuedByUser, nil, nil) + persistentA := repoA.(sessionTransferCompensationRepository) + if err := persistentA.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + serviceB := newAutoMergeTestServiceWithRepository(t, repoB, DefaultMaxPerSession) + + lease, err := serviceB.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + + if !errors.Is(err, ErrSessionTransferInProgress) { + t.Fatalf("BeginEdit error = %v, want %v", err, ErrSessionTransferInProgress) + } + if lease != nil { + t.Fatalf("lease = %#v, want nil", lease) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_postgres_merge_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_postgres_merge_test.go index 8c791ee6f8..3d33eaac2b 100644 --- a/apps/backend/internal/orchestrator/messagequeue/repository_postgres_merge_test.go +++ b/apps/backend/internal/orchestrator/messagequeue/repository_postgres_merge_test.go @@ -572,15 +572,16 @@ func TestPostgresRepository_TransferSession_RaceWithSourceInsert(t *testing.T) { } } -// TestPostgresRepository_AcknowledgeByID_RaceWithRestore proves the ack is -// cross-process serialized on the session lock: an autocommit DELETE could -// otherwise remove a row a concurrent backend just restored/reinserted, -// losing durable retry state. The competing backend holds the session lock, -// the ack blocks on it, and only after release does the ack delete the row. -func TestPostgresRepository_AcknowledgeByID_RaceWithSessionLock(t *testing.T) { +// TestPostgresRepository_AcknowledgeReserved_RaceWithSessionLock proves the +// attempt-scoped ack is cross-process serialized on the session lock. +func TestPostgresRepository_AcknowledgeReserved_RaceWithSessionLock(t *testing.T) { repoA, repoB, _ := newTestPostgresRepoPair(t) ctx := context.Background() - entry := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("ack", "first")) + entry := insertDurableLifecycleEntry(t, repoA, "ack") + reserved, err := repoA.ReserveHead(ctx, entry.SessionID) + if err != nil || reserved == nil { + t.Fatalf("reserve lifecycle entry: msg=%+v err=%v", reserved, err) + } dbA := repoA.(*sqliteRepository).db lockTx, err := dbA.BeginTxx(ctx, nil) @@ -603,7 +604,7 @@ func TestPostgresRepository_AcknowledgeByID_RaceWithSessionLock(t *testing.T) { ackPID := pgBackendPID(t, repoB.(*sqliteRepository).db) ackDone := make(chan error, 1) go func() { - ackDone <- repoB.AcknowledgeByID(ctx, "ack", entry.ID) + ackDone <- repoB.AcknowledgeReserved(ctx, reserved) }() waitForWaitingLocks(t, lockTx, ackPID, 1, "ack on the session lock") @@ -815,7 +816,7 @@ func TestPostgresRepository_CancellationIncludesAllOriginsAndPreservesReservatio repo := newTestPostgresRepo(t) ctx := context.Background() - reservedEntry := insertDurableLifecycleEntry(t, repo, "s1") + _ = insertDurableLifecycleEntry(t, repo, "s1") reserved, err := repo.ReserveHead(ctx, "s1") if err != nil || reserved == nil { t.Fatalf("reserve lifecycle entry: msg=%+v err=%v", reserved, err) @@ -835,7 +836,7 @@ func TestPostgresRepository_CancellationIncludesAllOriginsAndPreservesReservatio if removed != 4 { t.Fatalf("removed = %d, want 4", removed) } - if err := repo.AcknowledgeByID(ctx, "s1", reservedEntry.ID); err != nil { + if err := repo.AcknowledgeReserved(ctx, reserved); err != nil { t.Fatalf("reserved entry did not survive clear: %v", err) } } diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_send_now_recovery.go b/apps/backend/internal/orchestrator/messagequeue/repository_send_now_recovery.go new file mode 100644 index 0000000000..e734501d96 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_send_now_recovery.go @@ -0,0 +1,327 @@ +package messagequeue + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "github.com/google/uuid" + "github.com/jmoiron/sqlx" + internaldb "github.com/kandev/kandev/internal/db" +) + +const sendNowClaimRecoverySchema = ` + CREATE TABLE IF NOT EXISTS queue_send_now_claims ( + session_id TEXT PRIMARY KEY, + claim_id TEXT NOT NULL, + claim_json TEXT NOT NULL, + accepted INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMP NOT NULL + ) +` + +func (r *sqliteRepository) ensureSendNowClaimRecoverySchema(ctx context.Context) error { + if _, err := r.db.ExecContext(ctx, sendNowClaimRecoverySchema); err != nil { + return fmt.Errorf("ensure Send Now claim recovery schema: %w", err) + } + if _, err := r.db.ExecContext(ctx, `ALTER TABLE queue_send_now_claims ADD COLUMN accepted INTEGER NOT NULL DEFAULT 0`); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add Send Now claim acceptance: %w", err) + } + if _, err := r.db.ExecContext(ctx, `ALTER TABLE queue_send_now_claims ADD COLUMN claim_id TEXT NOT NULL DEFAULT ''`); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add Send Now claim identity: %w", err) + } + rows, err := r.db.QueryxContext(ctx, `SELECT session_id FROM queue_send_now_claims WHERE claim_id = ''`) + if err != nil { + return fmt.Errorf("list Send Now claims without identities: %w", err) + } + var sessionIDs []string + for rows.Next() { + var sessionID string + if err := rows.Scan(&sessionID); err != nil { + _ = rows.Close() + return fmt.Errorf("scan Send Now claim without identity: %w", err) + } + sessionIDs = append(sessionIDs, sessionID) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return fmt.Errorf("iterate Send Now claims without identities: %w", err) + } + if err := rows.Close(); err != nil { + return fmt.Errorf("close Send Now claim identity migration rows: %w", err) + } + for _, sessionID := range sessionIDs { + if _, err := r.db.ExecContext(ctx, r.db.Rebind(` + UPDATE queue_send_now_claims SET claim_id = ? + WHERE session_id = ? AND claim_id = '' + `), uuid.NewString(), sessionID); err != nil { + return fmt.Errorf("backfill Send Now claim identity: %w", err) + } + } + return nil +} + +func (r *sqliteRepository) persistSendNowClaimTx( + ctx context.Context, + tx *sqlx.Tx, + claim *SendNowClaim, +) error { + sessionID, err := sendNowClaimSessionID(claim) + if err != nil { + return err + } + if claim.ClaimID == "" { + claim.ClaimID = uuid.NewString() + } + claimJSON, err := json.Marshal(claim) + if err != nil { + return fmt.Errorf("marshal Send Now claim: %w", err) + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + INSERT INTO queue_send_now_claims (session_id, claim_id, claim_json, created_at) + VALUES (?, ?, ?, ?) + `), sessionID, claim.ClaimID, string(claimJSON), time.Now().UTC()); err != nil { + if isDuplicateSendNowClaimError(err) { + return ErrSendNowClaimChanged + } + return fmt.Errorf("persist Send Now claim: %w", err) + } + return nil +} + +func isDuplicateSendNowClaimError(err error) bool { + if err == nil { + return false + } + message := strings.ToLower(err.Error()) + return strings.Contains(message, "duplicate key") || + strings.Contains(message, "unique constraint failed: queue_send_now_claims.session_id") +} + +func (r *sqliteRepository) deleteSendNowClaimTx( + ctx context.Context, + tx *sqlx.Tx, + sessionID string, +) error { + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queue_send_now_claims WHERE session_id = ? + `), sessionID); err != nil { + return fmt.Errorf("delete Send Now claim: %w", err) + } + return nil +} + +func (r *sqliteRepository) deleteExactSendNowClaimTx( + ctx context.Context, + tx *sqlx.Tx, + sessionID, claimID string, +) error { + result, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queue_send_now_claims WHERE session_id = ? AND claim_id = ? + `), sessionID, claimID) + if err != nil { + return fmt.Errorf("delete exact Send Now claim: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("delete exact Send Now claim rows affected: %w", err) + } + if affected != 1 { + return ErrSendNowClaimChanged + } + return nil +} + +func (r *sqliteRepository) transferPendingSendNowClaimTx( + ctx context.Context, + tx *sqlx.Tx, + oldSessionID, newSessionID string, +) error { + var claimID, claimJSON string + err := tx.QueryRowxContext(ctx, r.db.Rebind(` + SELECT claim_id, claim_json + FROM queue_send_now_claims + WHERE session_id = ? + `), oldSessionID).Scan(&claimID, &claimJSON) + if errors.Is(err, sql.ErrNoRows) { + return nil + } + if err != nil { + return fmt.Errorf("read source Send Now claim: %w", err) + } + var destinationClaimID string + err = tx.QueryRowxContext(ctx, r.db.Rebind(` + SELECT claim_id + FROM queue_send_now_claims + WHERE session_id = ? + `), newSessionID).Scan(&destinationClaimID) + if err == nil { + return ErrSendNowClaimChanged + } + if !errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("read destination Send Now claim: %w", err) + } + var claim SendNowClaim + if err := json.Unmarshal([]byte(claimJSON), &claim); err != nil { + return fmt.Errorf("unmarshal source Send Now claim: %w", err) + } + claim.ClaimID = claimID + sessionID, err := sendNowClaimSessionID(&claim) + if err != nil || sessionID != oldSessionID || claim.Dispatch.SessionID != oldSessionID { + return ErrSendNowClaimChanged + } + for sourceIndex := range claim.Sources { + claim.Sources[sourceIndex].SessionID = newSessionID + } + claim.Dispatch.SessionID = newSessionID + claim.SessionGeneration, err = r.getSendNowGenerationTx(ctx, tx, newSessionID) + if err != nil { + return err + } + claimJSONBytes, err := json.Marshal(&claim) + if err != nil { + return fmt.Errorf("marshal transferred Send Now claim: %w", err) + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + UPDATE queue_send_now_claims + SET session_id = ?, claim_json = ? + WHERE session_id = ? AND claim_id = ? + `), newSessionID, string(claimJSONBytes), oldSessionID, claimID); err != nil { + return fmt.Errorf("transfer Send Now claim: %w", err) + } + return nil +} + +func (r *sqliteRepository) ListPendingSendNowClaims(ctx context.Context) ([]PendingSendNowClaim, error) { + if err := r.ensureSendNowClaimRecoverySchema(ctx); err != nil { + return nil, err + } + rows, err := r.db.QueryxContext( + ctx, + `SELECT claim_id, claim_json, accepted FROM queue_send_now_claims ORDER BY created_at, session_id`, + ) + if err != nil { + return nil, fmt.Errorf("list Send Now claims: %w", err) + } + defer func() { _ = rows.Close() }() + var pending []PendingSendNowClaim + for rows.Next() { + var claimID, claimJSON string + var accepted int + if err := rows.Scan(&claimID, &claimJSON, &accepted); err != nil { + return nil, fmt.Errorf("scan Send Now claim: %w", err) + } + var claim SendNowClaim + if err := json.Unmarshal([]byte(claimJSON), &claim); err != nil { + return nil, fmt.Errorf("unmarshal Send Now claim: %w", err) + } + claim.ClaimID = claimID + pending = append(pending, PendingSendNowClaim{Claim: claim, Accepted: accepted != 0}) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate Send Now claims: %w", err) + } + return pending, nil +} +func (r *sqliteRepository) MarkPendingSendNowClaimAccepted(ctx context.Context, claim *SendNowClaim) error { + if err := r.ensureSendNowClaimRecoverySchema(ctx); err != nil { + return err + } + sessionID, err := sendNowClaimSessionID(claim) + if err != nil { + return err + } + if claim.ClaimID == "" { + return ErrSendNowClaimChanged + } + tx, err := r.beginSessionMutationTx(ctx, sessionID, "mark Send Now claim accepted") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + result, err := tx.ExecContext(ctx, r.db.Rebind(` + UPDATE queue_send_now_claims SET accepted = 1 + WHERE session_id = ? AND claim_id = ? + `), sessionID, claim.ClaimID) + if err != nil { + return fmt.Errorf("mark Send Now claim accepted: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("mark Send Now claim accepted rows affected: %w", err) + } + if affected != 1 { + return ErrSendNowClaimChanged + } + return tx.Commit() +} + +func (r *sqliteRepository) DeletePendingSendNowClaim(ctx context.Context, claim *SendNowClaim) error { + if err := r.ensureSendNowClaimRecoverySchema(ctx); err != nil { + return err + } + sessionID, err := sendNowClaimSessionID(claim) + if err != nil { + return err + } + if claim.ClaimID == "" { + return ErrSendNowClaimChanged + } + tx, err := r.beginSessionMutationTx(ctx, sessionID, "discard pending Send Now claim") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + result, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queue_send_now_claims WHERE session_id = ? AND claim_id = ? + `), sessionID, claim.ClaimID) + if err != nil { + return fmt.Errorf("discard pending Send Now claim: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("discard pending Send Now claim rows affected: %w", err) + } + if affected != 1 { + return ErrSendNowClaimChanged + } + return tx.Commit() +} + +func pendingSendNowSessionsForTaskTx( + ctx context.Context, + tx *sqlx.Tx, + taskID string, +) ([]string, error) { + rows, err := tx.QueryxContext(ctx, `SELECT session_id, claim_json FROM queue_send_now_claims`) + if err != nil { + return nil, fmt.Errorf("list task Send Now claims: %w", err) + } + defer func() { _ = rows.Close() }() + var sessions []string + for rows.Next() { + var sessionID, claimJSON string + if err := rows.Scan(&sessionID, &claimJSON); err != nil { + return nil, fmt.Errorf("scan task Send Now claim: %w", err) + } + var claim SendNowClaim + if err := json.Unmarshal([]byte(claimJSON), &claim); err != nil { + return nil, fmt.Errorf("unmarshal task Send Now claim: %w", err) + } + matches := claim.Dispatch.TaskID == taskID + for i := range claim.Sources { + matches = matches || claim.Sources[i].TaskID == taskID + } + if matches { + sessions = append(sessions, sessionID) + } + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate task Send Now claims: %w", err) + } + return sessions, nil +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_send_now_recovery_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_send_now_recovery_test.go new file mode 100644 index 0000000000..c369a8bd82 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_send_now_recovery_test.go @@ -0,0 +1,318 @@ +package messagequeue + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "testing" +) + +func TestSQLiteSendNowClaimPersistsOrdinarySourceRecovery(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + source := insertTestEntry(t, repo, "session-1", "task-1", "ordinary prompt", QueuedByUser, nil, nil) + + claim, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*source}) + if err != nil { + t.Fatal(err) + } + if entries, listErr := repo.ListBySession(ctx, "session-1"); listErr != nil || len(entries) != 0 { + t.Fatalf("queue after claim = %#v, err=%v, want ordinary source removed", entries, listErr) + } + persistent, ok := repo.(interface { + ListPendingSendNowClaims(context.Context) ([]PendingSendNowClaim, error) + }) + if !ok { + t.Fatal("SQLite queue repository does not persist pending Send Now claims") + } + pending, err := persistent.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || pending[0].Claim.Dispatch.ID != claim.Dispatch.ID { + t.Fatalf("pending claims = %#v, want dispatch %s", pending, claim.Dispatch.ID) + } + + if err := repo.RestoreSendNowClaim(ctx, &pending[0].Claim); err != nil { + t.Fatal(err) + } + pending, err = persistent.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("pending claims after restore = %#v", pending) + } + entries, err := repo.ListBySession(ctx, "session-1") + + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 || entries[0].ID != source.ID { + t.Fatalf("restored queue = %#v, want source %s", entries, source.ID) + } +} +func TestSQLiteTransferSessionMovesPendingSendNowClaim(t *testing.T) { + for _, accepted := range []bool{false, true} { + t.Run(fmt.Sprintf("accepted=%t", accepted), func(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + source := insertTestEntry( + t, repo, "session-old", "task-1", "durable source", QueuedByWorkflow, nil, + map[string]interface{}{MetadataLifecycleDurable: true}, + ) + claim, err := repo.ClaimSendNow(ctx, source.SessionID, []QueuedMessage{*source}) + if err != nil { + t.Fatal(err) + } + if accepted { + if err := repo.MarkPendingSendNowClaimAccepted(ctx, claim); err != nil { + t.Fatal(err) + } + } + + if err := repo.TransferSession(ctx, "session-old", "session-new"); err != nil { + t.Fatal(err) + } + + pending, err := repo.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 { + t.Fatalf("pending claims = %#v, want one transferred claim", pending) + } + transferred := pending[0] + if transferred.Claim.Sources[0].SessionID != "session-new" || + transferred.Claim.Dispatch.SessionID != "session-new" { + t.Fatalf("transferred claim = %#v, want session-new", transferred.Claim) + } + if transferred.Accepted != accepted { + t.Fatalf("transferred accepted = %t, want %t", transferred.Accepted, accepted) + } + if accepted { + err = repo.AcknowledgeSendNowClaim(ctx, &transferred.Claim) + } else { + err = repo.RestoreSendNowClaim(ctx, &transferred.Claim) + } + if err != nil { + t.Fatal(err) + } + entries, err := repo.ListBySession(ctx, "session-new") + if err != nil { + t.Fatal(err) + } + if accepted && len(entries) != 0 { + t.Fatalf("acknowledged destination queue = %#v, want empty", entries) + } + if !accepted && (len(entries) != 1 || entries[0].ID != source.ID || entries[0].IsReservedInFlight()) { + t.Fatalf("restored destination queue = %#v", entries) + } + }) + } +} + +func TestSQLitePurgeRemovesDurableSendNowClaims(t *testing.T) { + for _, purge := range []struct { + name string + run func(context.Context, Repository) error + }{ + { + name: "session", + run: func(ctx context.Context, repo Repository) error { + _, err := repo.PurgeSession(ctx, "session-purge-send-now") + return err + }, + }, + { + name: "task", + run: func(ctx context.Context, repo Repository) error { + _, err := repo.PurgeTask(ctx, "task-purge-send-now") + return err + }, + }, + } { + for _, accepted := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/accepted=%t", purge.name, accepted), func(t *testing.T) { + repo := newTestSQLiteRepo(t) + ctx := context.Background() + entry := insertTestEntry( + t, repo, "session-purge-send-now", "task-purge-send-now", + "durable source", QueuedByUser, nil, nil, + ) + pendingClaim, err := repo.ClaimSendNow(ctx, entry.SessionID, []QueuedMessage{*entry}) + if err != nil { + t.Fatal(err) + } + persistent := repo.(pendingSendNowClaimRepository) + if accepted { + if err := persistent.MarkPendingSendNowClaimAccepted(ctx, pendingClaim); err != nil { + t.Fatal(err) + } + } + + if err := purge.run(ctx, repo); err != nil { + t.Fatal(err) + } + claims, err := persistent.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(claims) != 0 { + t.Fatalf("durable Send Now claims after purge = %#v, want none", claims) + } + }) + } + } +} + +func TestSQLiteSendNowAcknowledgeGenerationChangeRetiresClaim(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + first := insertTestEntry(t, repo, "session-1", "task-1", "first", QueuedByUser, nil, nil) + claim, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*first}) + if err != nil { + t.Fatal(err) + } + if _, err := repo.DeleteAllBySession(ctx, "session-1"); err != nil { + t.Fatal(err) + } + if err := repo.AcknowledgeSendNowClaim(ctx, claim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("acknowledge after generation change error = %v, want %v", err, ErrSendNowClaimChanged) + } + + second := insertTestEntry(t, repo, "session-1", "task-1", "second", QueuedByUser, nil, nil) + if _, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*second}); err != nil { + t.Fatalf("second same-process Send Now claim: %v", err) + } +} + +func TestSQLiteStaleSendNowWorkerCannotSettleSuccessorClaim(t *testing.T) { + for _, operation := range []struct { + name string + run func(context.Context, *sqliteRepository, *SendNowClaim) error + }{ + { + name: "restore", + run: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.RestoreSendNowClaim(ctx, claim) + }, + }, + { + name: "acknowledge", + run: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.AcknowledgeSendNowClaim(ctx, claim) + }, + }, + { + name: "acknowledge-unidentified", + run: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + unidentified := *claim + unidentified.ClaimID = "" + return repo.AcknowledgeSendNowClaim(ctx, &unidentified) + }, + }, + { + name: "mark-accepted", + run: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.MarkPendingSendNowClaimAccepted(ctx, claim) + }, + }, + { + name: "startup-discard", + run: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.DeletePendingSendNowClaim(ctx, claim) + }, + }, + } { + t.Run(operation.name, func(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + oldSource := insertTestEntry(t, repo, "session-stale-worker", "task-stale-worker", "old", QueuedByUser, nil, nil) + oldClaim, err := repo.ClaimSendNow(ctx, oldSource.SessionID, []QueuedMessage{*oldSource}) + if err != nil { + t.Fatal(err) + } + if err := repo.AcknowledgeSendNowClaim(ctx, oldClaim); err != nil { + t.Fatal(err) + } + successorSource := insertTestEntry( + t, repo, oldSource.SessionID, oldSource.TaskID, "successor", QueuedByUser, nil, nil, + ) + successor, err := repo.ClaimSendNow(ctx, successorSource.SessionID, []QueuedMessage{*successorSource}) + if err != nil { + t.Fatal(err) + } + + if err := operation.run(ctx, repo, oldClaim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("stale %s error = %v, want %v", operation.name, err, ErrSendNowClaimChanged) + } + pending, err := repo.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || pending[0].Claim.Dispatch.ID != successor.Dispatch.ID || pending[0].Accepted { + t.Fatalf("successor after stale %s = %#v", operation.name, pending) + } + entries, err := repo.ListBySession(ctx, oldSource.SessionID) + if err != nil { + t.Fatal(err) + } + if len(entries) != 0 { + t.Fatalf("queue after stale %s = %#v, want successor sources still claimed", operation.name, entries) + } + }) + } +} + +func TestSQLiteSendNowClaimMigratesLegacyIdentity(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + _, err := repo.db.ExecContext(ctx, ` + CREATE TABLE queue_send_now_claims ( + session_id TEXT PRIMARY KEY, + claim_json TEXT NOT NULL, + accepted INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMP NOT NULL + ) + `) + if err != nil { + t.Fatal(err) + } + legacy := SendNowClaim{ + Sources: []QueuedMessage{{ + ID: "legacy-source", SessionID: "legacy-session", TaskID: "legacy-task", + Content: "legacy", QueuedBy: QueuedByUser, + }}, + Dispatch: QueuedMessage{ID: "legacy-dispatch", SessionID: "legacy-session", TaskID: "legacy-task"}, + } + claimJSON, err := json.Marshal(legacy) + if err != nil { + t.Fatal(err) + } + if _, err := repo.db.ExecContext(ctx, ` + INSERT INTO queue_send_now_claims (session_id, claim_json, accepted, created_at) + VALUES (?, ?, 0, CURRENT_TIMESTAMP) + `, legacy.Dispatch.SessionID, string(claimJSON)); err != nil { + t.Fatal(err) + } + + pending, err := repo.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || pending[0].Claim.ClaimID == "" { + t.Fatalf("migrated Send Now claim = %#v", pending) + } + if err := repo.MarkPendingSendNowClaimAccepted(ctx, &pending[0].Claim); err != nil { + t.Fatal(err) + } + if err := repo.DeletePendingSendNowClaim(ctx, &pending[0].Claim); err != nil { + t.Fatal(err) + } + pending, err = repo.ListPendingSendNowClaims(ctx) + if err != nil || len(pending) != 0 { + t.Fatalf("pending claims after migrated settlement = %#v, err=%v", pending, err) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_send_now_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_send_now_test.go index 7f1d8fa7c9..677fd848ee 100644 --- a/apps/backend/internal/orchestrator/messagequeue/repository_send_now_test.go +++ b/apps/backend/internal/orchestrator/messagequeue/repository_send_now_test.go @@ -4,8 +4,43 @@ import ( "context" "errors" "testing" + + "github.com/stretchr/testify/require" ) +func TestSessionGenerationAdvancesForDestructiveQueueChanges(t *testing.T) { + tests := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := tt.new(t) + ctx := context.Background() + if got, err := repo.SessionGeneration(ctx, "session-1"); err != nil || got != 0 { + t.Fatalf("initial session generation = %d err=%v, want 0", got, err) + } + insertTestEntry(t, repo, "session-1", "task-1", "pending", QueuedByUser, nil, nil) + if _, err := repo.DeleteAllBySession(ctx, "session-1"); err != nil { + t.Fatalf("delete session queue: %v", err) + } + if got, err := repo.SessionGeneration(ctx, "session-1"); err != nil || got != 1 { + t.Fatalf("session generation after delete = %d err=%v, want 1", got, err) + } + if _, err := repo.PurgeSession(ctx, "session-1"); err != nil { + t.Fatalf("purge session queue: %v", err) + } + if got, err := repo.SessionGeneration(ctx, "session-1"); err != nil || got != 2 { + t.Fatalf("session generation after purge = %d err=%v, want 2", got, err) + } + }) + } +} + func TestSendNowClaimIsExactAtomicAndRestorable(t *testing.T) { tests := []struct { name string @@ -428,7 +463,8 @@ func TestSendNowRestoreDiscardsSourcesFromPurgedTaskGeneration(t *testing.T) { if _, err := repo.PurgeTask(ctx, "task-1"); err != nil { t.Fatalf("purge task: %v", err) } - if err := repo.RestoreSendNowClaim(ctx, claim); err != nil { + err = repo.RestoreSendNowClaim(ctx, claim) + if err != nil { t.Fatalf("restore after purge: %v", err) } entries, err := repo.ListBySession(ctx, "session-1") @@ -441,7 +477,241 @@ func TestSendNowRestoreDiscardsSourcesFromPurgedTaskGeneration(t *testing.T) { }) } } +func TestSendNowRestoreDoesNotResurrectSourcesFromPurgedSession(t *testing.T) { + tests := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := tt.new(t) + ctx := context.Background() + entry := insertTestEntry(t, repo, "session-1", "task-1", "discard after session purge", QueuedByUser, nil, nil) + claim, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*entry}) + if err != nil { + t.Fatalf("claim: %v", err) + } + if _, err := repo.PurgeSession(ctx, "session-1"); err != nil { + t.Fatalf("purge session: %v", err) + } + if err := repo.RestoreSendNowClaim(ctx, claim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("restore after session purge error = %v, want ErrSendNowClaimChanged", err) + } + entries, err := repo.ListBySession(ctx, "session-1") + if err != nil { + t.Fatalf("list after restore: %v", err) + } + if len(entries) != 0 { + t.Fatalf("restored entries after session purge = %#v, want none", entries) + } + }) + } +} + +func TestSendNowRestoreFencesSessionTransferAndReplacement(t *testing.T) { + tests := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Run("transfer", func(t *testing.T) { + repo := tt.new(t) + ctx := context.Background() + entry := insertTestEntry(t, repo, "session-1", "task-1", "transferred", QueuedByUser, nil, nil) + claim, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*entry}) + if err != nil { + t.Fatalf("claim: %v", err) + } + if err := repo.TransferSession(ctx, "session-1", "session-2"); err != nil { + t.Fatalf("transfer: %v", err) + } + if err := repo.RestoreSendNowClaim(ctx, claim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("restore after transfer error = %v, want ErrSendNowClaimChanged", err) + } + entries, err := repo.ListBySession(ctx, "session-2") + if err != nil { + t.Fatalf("list transferred entries: %v", err) + } + if len(entries) != 0 { + t.Fatalf("transferred entries = %#v, want none", entries) + } + }) + + t.Run("replacement", func(t *testing.T) { + repo := tt.new(t) + ctx := context.Background() + entry := insertTestEntry(t, repo, "session-1", "task-1", "replaced", QueuedByUser, nil, nil) + claim, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*entry}) + if err != nil { + t.Fatalf("claim: %v", err) + } + if err := repo.ReplaceSession(ctx, "session-1", nil, nil); err != nil { + t.Fatalf("replace: %v", err) + } + if err := repo.RestoreSendNowClaim(ctx, claim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("restore after replacement error = %v, want ErrSendNowClaimChanged", err) + } + entries, err := repo.ListBySession(ctx, "session-1") + if err != nil { + t.Fatalf("list replaced entries: %v", err) + } + if len(entries) != 0 { + t.Fatalf("replaced entries = %#v, want none", entries) + } + }) + }) + } +} + +func TestSendNowAcknowledgeFencesSessionChanges(t *testing.T) { + tests := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for _, change := range []struct { + name string + run func(*testing.T, Repository, *SendNowClaim) + }{ + { + name: "transfer", + run: func(t *testing.T, repo Repository, claim *SendNowClaim) { + ctx := context.Background() + if err := repo.TransferSession(ctx, "session-1", "session-2"); err != nil { + t.Fatalf("transfer: %v", err) + } + if err := repo.AcknowledgeSendNowClaim(ctx, claim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("acknowledge after transfer error = %v, want ErrSendNowClaimChanged", err) + } + entries, err := repo.ListBySession(ctx, "session-2") + if err != nil { + t.Fatalf("list transferred entries: %v", err) + } + if len(entries) != 1 || !entries[0].IsReservedInFlight() { + t.Fatalf("transferred entries = %#v, want reserved source preserved", entries) + } + }, + }, + { + name: "replacement", + run: func(t *testing.T, repo Repository, claim *SendNowClaim) { + ctx := context.Background() + replacement, err := repo.ListBySession(ctx, "session-1") + if err != nil { + t.Fatalf("list reserved replacement: %v", err) + } + if err := repo.ReplaceSession(ctx, "session-1", replacement, nil); err != nil { + t.Fatalf("replace: %v", err) + } + if err := repo.AcknowledgeSendNowClaim(ctx, claim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("acknowledge after replacement error = %v, want ErrSendNowClaimChanged", err) + } + entries, err := repo.ListBySession(ctx, "session-1") + if err != nil { + t.Fatalf("list replaced entries: %v", err) + } + if len(entries) != 1 || !entries[0].IsReservedInFlight() { + t.Fatalf("replaced entries = %#v, want reserved replacement preserved", entries) + } + }, + }, + { + name: "purge", + run: func(t *testing.T, repo Repository, claim *SendNowClaim) { + ctx := context.Background() + if _, err := repo.PurgeSession(ctx, "session-1"); err != nil { + t.Fatalf("purge session: %v", err) + } + if err := repo.AcknowledgeSendNowClaim(ctx, claim); !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("acknowledge after purge error = %v, want ErrSendNowClaimChanged", err) + } + }, + }, + } { + t.Run(change.name, func(t *testing.T) { + repo := tt.new(t) + ctx := context.Background() + entry := insertTestEntry(t, repo, "session-1", "task-1", "durable source", QueuedByWorkflow, nil, + map[string]interface{}{MetadataLifecycleDurable: true}) + claim, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*entry}) + if err != nil { + t.Fatalf("claim: %v", err) + } + change.run(t, repo, claim) + }) + } + }) + } +} + +func TestSendNowAcknowledgeIgnoresPurgedTaskGeneration(t *testing.T) { + tests := []struct { + name string + new func(*testing.T) Repository + }{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := tt.new(t) + ctx := context.Background() + entry := insertTestEntry(t, repo, "session-1", "task-1", "durable source", QueuedByWorkflow, nil, + map[string]interface{}{MetadataLifecycleDurable: true}) + claim, err := repo.ClaimSendNow(ctx, "session-1", []QueuedMessage{*entry}) + if err != nil { + t.Fatalf("claim: %v", err) + } + if _, err := repo.PurgeTask(ctx, "task-1"); err != nil { + t.Fatalf("purge task: %v", err) + } + + replacement := claim.Sources[0] + replacement.Metadata = map[string]interface{}{ + MetadataLifecycleDurable: true, + MetadataLifecycleReserved: true, + } + replacement.reservedLifecycleDelivery = true + replacement.SessionID = "session-1" + if err := repo.ReplaceSession(ctx, "session-1", []QueuedMessage{replacement}, nil); err != nil { + t.Fatalf("replace with successor: %v", err) + } + // The replacement changes the session generation, so isolate the + // task-generation fence to the purged task generation itself. + claim.SessionGeneration++ + err = repo.AcknowledgeSendNowClaim(ctx, claim) + if tt.name == "sqlite" { + if !errors.Is(err, ErrSendNowClaimChanged) { + t.Fatalf("acknowledge stale task claim error = %v, want %v", err, ErrSendNowClaimChanged) + } + } else if err != nil { + t.Fatalf("acknowledge stale task claim: %v", err) + } + entries, err := repo.ListBySession(ctx, "session-1") + if err != nil { + t.Fatalf("list successor: %v", err) + } + if len(entries) != 1 || !entries[0].IsReservedInFlight() { + t.Fatalf("successor after stale acknowledge = %#v, want reserved successor preserved", entries) + } + }) + } +} func TestSendNowRestoreKeepsCurrentDurableMetadataAndRecordedMarker(t *testing.T) { repo := NewMemoryRepository().(*memoryRepository) ctx := context.Background() @@ -476,3 +746,101 @@ func TestSendNowRestoreKeepsCurrentDurableMetadataAndRecordedMarker(t *testing.T t.Fatalf("restored metadata lost recorded marker: %#v", entries[0].Metadata) } } +func TestRestoreKeepsPositionsUniqueAfterQueueWasEmptied(t *testing.T) { + for _, factory := range autoRunRepositoryFactories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + ctx := context.Background() + first := insertTestEntry(t, repo, "session-restore-order", "task", "first", QueuedByUser, nil, nil) + second := insertTestEntry(t, repo, "session-restore-order", "task", "second", QueuedByUser, nil, nil) + sources := []QueuedMessage{*first, *second} + claim, err := repo.ClaimSendNow(ctx, "session-restore-order", sources) + require.NoError(t, err) + later := insertTestEntry(t, repo, "session-restore-order", "task", "later", QueuedByUser, nil, nil) + require.NoError(t, repo.RestoreSendNowClaim(ctx, claim)) + + entries, err := repo.ListBySession(ctx, "session-restore-order") + require.NoError(t, err) + require.Len(t, entries, 3) + require.Equal(t, []string{first.ID, second.ID, later.ID}, []string{entries[0].ID, entries[1].ID, entries[2].ID}) + require.Less(t, entries[0].Position, entries[1].Position) + require.Less(t, entries[1].Position, entries[2].Position) + }) + } +} +func TestOrdinaryRestoreKeepsPositionAheadOfLaterArrival(t *testing.T) { + for _, factory := range autoRunRepositoryFactories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + ctx := context.Background() + first := insertTestEntry(t, repo, "session-ordinary-restore", "task", "first", QueuedByUser, nil, nil) + later := insertTestEntry(t, repo, "session-ordinary-restore", "task", "later", QueuedByUser, nil, nil) + taken, err := repo.TakeHead(ctx, "session-ordinary-restore") + require.NoError(t, err) + require.Equal(t, first.ID, taken.ID) + third := insertTestEntry(t, repo, "session-ordinary-restore", "task", "third", QueuedByUser, nil, nil) + require.NoError(t, repo.Restore(ctx, taken, 0)) + + entries, err := repo.ListBySession(ctx, "session-ordinary-restore") + require.NoError(t, err) + require.Len(t, entries, 3) + require.Equal(t, []string{first.ID, later.ID, third.ID}, []string{entries[0].ID, entries[1].ID, entries[2].ID}) + require.Less(t, entries[0].Position, entries[1].Position) + require.Less(t, entries[1].Position, entries[2].Position) + }) + } +} +func TestTransferRestoreKeepsDestinationPositionHighWater(t *testing.T) { + for _, factory := range autoRunRepositoryFactories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + ctx := context.Background() + var transferredSource *QueuedMessage + for i := range 7 { + entry := insertTestEntry(t, repo, "transfer-source", "task", "source", QueuedByUser, nil, nil) + if i == 6 { + transferredSource = entry + continue + } + taken, err := repo.TakeHead(ctx, "transfer-source") + require.NoError(t, err) + require.Equal(t, entry.ID, taken.ID) + } + require.NotNil(t, transferredSource) + + require.NoError(t, repo.TransferSession(ctx, "transfer-source", "transfer-destination")) + transferred, err := repo.TakeHead(ctx, "transfer-destination") + require.NoError(t, err) + require.Equal(t, transferredSource.ID, transferred.ID) + later := insertTestEntry(t, repo, "transfer-destination", "task", "later", QueuedByUser, nil, nil) + + require.NoError(t, repo.Restore(ctx, transferred, 0)) + entries, err := repo.ListBySession(ctx, "transfer-destination") + require.NoError(t, err) + require.Equal(t, []string{transferred.ID, later.ID}, []string{entries[0].ID, entries[1].ID}) + require.Less(t, entries[0].Position, entries[1].Position) + }) + } +} +func TestOrdinaryRestoreKeepsPositionAfterDeleteEmptiesQueue(t *testing.T) { + for _, factory := range autoRunRepositoryFactories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + ctx := context.Background() + first := insertTestEntry(t, repo, "session-delete-restore", "task", "first", QueuedByUser, nil, nil) + second := insertTestEntry(t, repo, "session-delete-restore", "task", "second", QueuedByUser, nil, nil) + taken, err := repo.TakeHead(ctx, "session-delete-restore") + require.NoError(t, err) + require.Equal(t, first.ID, taken.ID) + require.NoError(t, repo.DeleteByID(ctx, "session-delete-restore", second.ID)) + + later := insertTestEntry(t, repo, "session-delete-restore", "task", "later", QueuedByUser, nil, nil) + require.NoError(t, repo.Restore(ctx, taken, 0)) + + entries, err := repo.ListBySession(ctx, "session-delete-restore") + require.NoError(t, err) + require.Equal(t, []string{first.ID, later.ID}, []string{entries[0].ID, entries[1].ID}) + require.Less(t, entries[0].Position, entries[1].Position) + }) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_session_transfer_compensation.go b/apps/backend/internal/orchestrator/messagequeue/repository_session_transfer_compensation.go new file mode 100644 index 0000000000..ee92ce1d65 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_session_transfer_compensation.go @@ -0,0 +1,559 @@ +package messagequeue + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "github.com/google/uuid" + "github.com/jmoiron/sqlx" + internaldb "github.com/kandev/kandev/internal/db" + "time" +) + +const sessionTransferCompensationLeaseDuration = time.Minute +const postgresDriverName = "pgx" +const postgresForUpdateSuffix = " FOR UPDATE" + +const sessionTransferCompensationSchema = ` + CREATE TABLE IF NOT EXISTS queue_session_transfer_compensations ( + task_id TEXT NOT NULL, + from_session_id TEXT NOT NULL, + to_session_id TEXT NOT NULL, + operation_id TEXT NOT NULL DEFAULT '', + recovery_owner TEXT NOT NULL DEFAULT '', + recovery_lease_expires_at TIMESTAMP, + entry_ids_json TEXT NOT NULL DEFAULT '[]', + attachment_ids_json TEXT NOT NULL DEFAULT '[]', + cleanup_locators_json TEXT NOT NULL DEFAULT '[]', + created_at TIMESTAMP NOT NULL, + PRIMARY KEY (task_id, from_session_id, to_session_id) + ) +` + +func (r *sqliteRepository) ensureSessionTransferCompensationSchema(ctx context.Context) error { + if _, err := r.db.ExecContext(ctx, sessionTransferCompensationSchema); err != nil { + return fmt.Errorf("ensure session transfer compensation schema: %w", err) + } + for _, migration := range []struct { + statement string + name string + }{ + {`ALTER TABLE queue_session_transfer_compensations ADD COLUMN operation_id TEXT NOT NULL DEFAULT ''`, "operation id"}, + {`ALTER TABLE queue_session_transfer_compensations ADD COLUMN attachment_ids_json TEXT NOT NULL DEFAULT '[]'`, "attachment ids"}, + {`ALTER TABLE queue_session_transfer_compensations ADD COLUMN cleanup_locators_json TEXT NOT NULL DEFAULT '[]'`, "cleanup locators"}, + {`ALTER TABLE queue_session_transfer_compensations ADD COLUMN recovery_owner TEXT NOT NULL DEFAULT ''`, "recovery owner"}, + {`ALTER TABLE queue_session_transfer_compensations ADD COLUMN recovery_lease_expires_at TIMESTAMP`, "recovery lease"}, + } { + if _, err := r.db.ExecContext(ctx, migration.statement); err != nil && !internaldb.IsDuplicateColumnError(err) { + return fmt.Errorf("add session transfer %s: %w", migration.name, err) + } + } + if _, err := r.db.ExecContext(ctx, ` + UPDATE queue_session_transfer_compensations + SET recovery_owner = operation_id + WHERE recovery_owner = '' + `); err != nil { + return fmt.Errorf("backfill session transfer recovery owner: %w", err) + } + return nil +} + +func marshalSessionTransferCompensation( + compensation SessionTransferCompensation, +) (string, string, string, error) { + entryIDsJSON, err := json.Marshal(compensation.EntryIDs) + if err != nil { + return "", "", "", fmt.Errorf("marshal session transfer compensation: %w", err) + } + attachmentIDsJSON, err := json.Marshal(compensation.AttachmentIDs) + if err != nil { + return "", "", "", fmt.Errorf("marshal session transfer attachment ids: %w", err) + } + cleanupLocatorsJSON, err := json.Marshal(compensation.CleanupLocators) + if err != nil { + return "", "", "", fmt.Errorf("marshal session transfer cleanup locators: %w", err) + } + return string(entryIDsJSON), string(attachmentIDsJSON), string(cleanupLocatorsJSON), nil +} + +func (r *sqliteRepository) lockSessionTransferPairTx( + ctx context.Context, + tx *sqlx.Tx, + fromSessionID, toSessionID string, +) (string, string, error) { + first, second := fromSessionID, toSessionID + if first > second { + first, second = second, first + } + if err := r.lockSessionTxUnfenced(ctx, tx, first); err != nil { + return "", "", err + } + if first != second { + if err := r.lockSessionTxUnfenced(ctx, tx, second); err != nil { + return "", "", err + } + } + return first, second, nil +} + +func guardSessionTransferCompensationOwnerTx( + ctx context.Context, + tx *sqlx.Tx, + compensation SessionTransferCompensation, + first, second string, +) error { + var conflicts int + if err := tx.GetContext(ctx, &conflicts, tx.Rebind(` + SELECT COUNT(*) FROM queue_session_transfer_compensations + WHERE (from_session_id IN (?, ?) OR to_session_id IN (?, ?)) + AND NOT (task_id = ? AND from_session_id = ? AND to_session_id = ?) + `), first, second, first, second, compensation.TaskID, + compensation.FromSessionID, compensation.ToSessionID); err != nil { + return fmt.Errorf("check active session transfer: %w", err) + } + if conflicts > 0 { + return ErrSessionTransferInProgress + } + var activeOperationID string + err := tx.GetContext(ctx, &activeOperationID, tx.Rebind(` + SELECT operation_id FROM queue_session_transfer_compensations + WHERE task_id = ? AND from_session_id = ? AND to_session_id = ? + `), compensation.TaskID, compensation.FromSessionID, compensation.ToSessionID) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("read active session transfer owner: %w", err) + } + if err == nil && activeOperationID != compensation.OperationID { + return ErrSessionTransferInProgress + } + return nil +} + +func (r *sqliteRepository) lockTaskForSessionTransferTx( + ctx context.Context, + tx *sqlx.Tx, + taskID string, + fromSessionID string, + toSessionID string, +) error { + if !r.tasksTablePresent { + return nil + } + query := `SELECT id FROM tasks WHERE id = ?` + if r.db.DriverName() == postgresDriverName { + query += postgresForUpdateSuffix + } + var lockedTaskID string + if err := tx.GetContext(ctx, &lockedTaskID, r.db.Rebind(query), taskID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("session transfer task %q no longer exists", taskID) + } + return fmt.Errorf("lock session transfer task %q: %w", taskID, err) + } + if !r.taskSessionsTablePresent { + return nil + } + for _, sessionID := range []string{fromSessionID, toSessionID} { + if sessionID == "" { + continue + } + var sessionTaskID string + if err := tx.GetContext(ctx, &sessionTaskID, r.db.Rebind(` + SELECT task_id FROM task_sessions WHERE id = ? + `), sessionID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("session transfer session %q no longer exists", sessionID) + } + return fmt.Errorf("load session transfer session %q: %w", sessionID, err) + } + if sessionTaskID != taskID { + return fmt.Errorf("session transfer session %q belongs to task %q, not %q", sessionID, sessionTaskID, taskID) + } + } + return nil +} + +func (r *sqliteRepository) UpsertSessionTransferCompensation( + ctx context.Context, + compensation SessionTransferCompensation, +) error { + if compensation.OperationID == "" { + return errors.New("session transfer operation id is required") + } + if err := r.ensureSessionTransferCompensationSchema(ctx); err != nil { + return err + } + entryIDsJSON, attachmentIDsJSON, cleanupLocatorsJSON, err := marshalSessionTransferCompensation(compensation) + if err != nil { + return err + } + if compensation.CreatedAt.IsZero() { + compensation.CreatedAt = time.Now().UTC() + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin session transfer compensation: %w", err) + } + defer func() { _ = tx.Rollback() }() + if err := r.lockTaskForSessionTransferTx( + ctx, tx, compensation.TaskID, compensation.FromSessionID, compensation.ToSessionID, + ); err != nil { + return err + } + first, second, err := r.lockSessionTransferPairTx( + ctx, tx, compensation.FromSessionID, compensation.ToSessionID, + ) + if err != nil { + return err + } + if err := guardSessionTransferCompensationOwnerTx( + ctx, tx, compensation, first, second, + ); err != nil { + return err + } + leaseExpiresAt := time.Now().UTC().Add(sessionTransferCompensationLeaseDuration) + result, err := tx.ExecContext(ctx, tx.Rebind(` + INSERT INTO queue_session_transfer_compensations + (operation_id, recovery_owner, recovery_lease_expires_at, task_id, + from_session_id, to_session_id, entry_ids_json, attachment_ids_json, + cleanup_locators_json, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(task_id, from_session_id, to_session_id) DO UPDATE SET + entry_ids_json = excluded.entry_ids_json, + attachment_ids_json = excluded.attachment_ids_json, + cleanup_locators_json = excluded.cleanup_locators_json, + recovery_lease_expires_at = excluded.recovery_lease_expires_at + WHERE queue_session_transfer_compensations.operation_id = excluded.operation_id + AND queue_session_transfer_compensations.recovery_owner = excluded.recovery_owner + `), compensation.OperationID, compensation.OperationID, leaseExpiresAt, + compensation.TaskID, compensation.FromSessionID, compensation.ToSessionID, + entryIDsJSON, attachmentIDsJSON, cleanupLocatorsJSON, compensation.CreatedAt) + if err != nil { + return fmt.Errorf("upsert session transfer compensation: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("session transfer compensation rows affected: %w", err) + } + if affected != 1 { + return ErrSessionTransferOwnershipLost + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit session transfer compensation: %w", err) + } + return nil +} + +func (r *sqliteRepository) DeleteSessionTransferCompensation( + ctx context.Context, + operationID, taskID, fromSessionID, toSessionID string, +) error { + return r.deleteSessionTransferCompensationWithOwner( + ctx, operationID, operationID, taskID, fromSessionID, toSessionID, + ) +} + +func (r *sqliteRepository) deleteSessionTransferCompensationWithOwner( + ctx context.Context, + operationID, ownerID, taskID, fromSessionID, toSessionID string, +) error { + if err := r.ensureSessionTransferCompensationSchema(ctx); err != nil { + return err + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin delete session transfer compensation: %w", err) + } + defer func() { _ = tx.Rollback() }() + if _, _, err := r.lockSessionTransferPairTx( + ctx, tx, fromSessionID, toSessionID, + ); err != nil { + return err + } + now := time.Now().UTC() + result, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM queue_session_transfer_compensations + WHERE operation_id = ? AND recovery_owner = ? + AND task_id = ? AND from_session_id = ? AND to_session_id = ? + AND recovery_lease_expires_at > ? + `), operationID, ownerID, taskID, fromSessionID, toSessionID, now) + if err != nil { + return fmt.Errorf("delete session transfer compensation: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("delete session transfer compensation rows affected: %w", err) + } + if affected != 1 { + return ErrSessionTransferOwnershipLost + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit delete session transfer compensation: %w", err) + } + return nil +} + +func (r *sqliteRepository) claimSessionTransferCompensationRecovery( + ctx context.Context, + compensation SessionTransferCompensation, +) (string, error) { + if err := r.ensureSessionTransferCompensationSchema(ctx); err != nil { + return "", err + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return "", fmt.Errorf("begin claim session transfer compensation: %w", err) + } + defer func() { _ = tx.Rollback() }() + if _, _, err := r.lockSessionTransferPairTx( + ctx, tx, compensation.FromSessionID, compensation.ToSessionID, + ); err != nil { + return "", err + } + ownerID := uuid.NewString() + now := time.Now().UTC() + result, err := tx.ExecContext(ctx, tx.Rebind(` + UPDATE queue_session_transfer_compensations + SET recovery_owner = ?, recovery_lease_expires_at = ? + WHERE operation_id = ? AND task_id = ? AND from_session_id = ? AND to_session_id = ? + AND (recovery_lease_expires_at IS NULL OR recovery_lease_expires_at <= ?) + `), ownerID, now.Add(sessionTransferCompensationLeaseDuration), + compensation.OperationID, compensation.TaskID, compensation.FromSessionID, + compensation.ToSessionID, now) + if err != nil { + return "", fmt.Errorf("claim session transfer compensation: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return "", fmt.Errorf("claim session transfer compensation rows affected: %w", err) + } + if affected != 1 { + return "", ErrSessionTransferInProgress + } + if err := tx.Commit(); err != nil { + return "", fmt.Errorf("commit claim session transfer compensation: %w", err) + } + return ownerID, nil +} + +func (r *sqliteRepository) renewSessionTransferCompensationLease( + ctx context.Context, + compensation SessionTransferCompensation, + ownerID string, +) error { + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin renew session transfer compensation: %w", err) + } + defer func() { _ = tx.Rollback() }() + if _, _, err := r.lockSessionTransferPairTx( + ctx, tx, compensation.FromSessionID, compensation.ToSessionID, + ); err != nil { + return err + } + now := time.Now().UTC() + result, err := tx.ExecContext(ctx, tx.Rebind(` + UPDATE queue_session_transfer_compensations + SET recovery_lease_expires_at = ? + WHERE operation_id = ? AND recovery_owner = ? + AND task_id = ? AND from_session_id = ? AND to_session_id = ? + AND recovery_lease_expires_at > ? + `), now.Add(sessionTransferCompensationLeaseDuration), + compensation.OperationID, ownerID, compensation.TaskID, + compensation.FromSessionID, compensation.ToSessionID, now) + if err != nil { + return fmt.Errorf("renew session transfer compensation: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("renew session transfer compensation rows affected: %w", err) + } + if affected != 1 { + return ErrSessionTransferOwnershipLost + } + return tx.Commit() +} + +func (r *sqliteRepository) ListSessionTransferCompensations( + ctx context.Context, +) ([]SessionTransferCompensation, error) { + if err := r.ensureSessionTransferCompensationSchema(ctx); err != nil { + return nil, err + } + rows, err := r.db.QueryxContext(ctx, ` + SELECT operation_id, task_id, from_session_id, to_session_id, entry_ids_json, + attachment_ids_json, cleanup_locators_json, created_at + FROM queue_session_transfer_compensations + ORDER BY created_at, task_id, from_session_id, to_session_id + `) + if err != nil { + return nil, fmt.Errorf("list session transfer compensations: %w", err) + } + defer func() { _ = rows.Close() }() + var compensations []SessionTransferCompensation + for rows.Next() { + var compensation SessionTransferCompensation + var entryIDsJSON, attachmentIDsJSON, cleanupLocatorsJSON string + if err := rows.Scan( + &compensation.OperationID, + &compensation.TaskID, + &compensation.FromSessionID, + &compensation.ToSessionID, + &entryIDsJSON, + &attachmentIDsJSON, + &cleanupLocatorsJSON, + &compensation.CreatedAt, + ); err != nil { + return nil, fmt.Errorf("scan session transfer compensation: %w", err) + } + if err := json.Unmarshal([]byte(entryIDsJSON), &compensation.EntryIDs); err != nil { + return nil, fmt.Errorf("unmarshal session transfer compensation: %w", err) + } + if err := json.Unmarshal([]byte(attachmentIDsJSON), &compensation.AttachmentIDs); err != nil { + return nil, fmt.Errorf("unmarshal session transfer attachment ids: %w", err) + } + if err := json.Unmarshal([]byte(cleanupLocatorsJSON), &compensation.CleanupLocators); err != nil { + return nil, fmt.Errorf("unmarshal session transfer cleanup locators: %w", err) + } + compensations = append(compensations, compensation) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate session transfer compensations: %w", err) + } + return compensations, nil +} + +func authorizeSessionTransferTx( + ctx context.Context, + tx *sqlx.Tx, + db *sqlx.DB, + fromSessionID, toSessionID, operationID string, +) error { + var activeOperationID, activeOwnerID, activeFromSessionID, activeToSessionID string + var leaseExpiresAt sql.NullTime + rows, err := tx.QueryxContext(ctx, db.Rebind(` + SELECT operation_id, recovery_owner, recovery_lease_expires_at, + from_session_id, to_session_id + FROM queue_session_transfer_compensations + WHERE from_session_id IN (?, ?) OR to_session_id IN (?, ?) + `), fromSessionID, toSessionID, fromSessionID, toSessionID) + if err != nil { + return fmt.Errorf("authorize session transfer: %w", err) + } + defer func() { _ = rows.Close() }() + if !rows.Next() { + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate session transfer authorization: %w", err) + } + if operationID != "" { + return ErrSessionTransferOwnershipLost + } + return nil + } + if err := rows.Scan( + &activeOperationID, &activeOwnerID, &leaseExpiresAt, + &activeFromSessionID, &activeToSessionID, + ); err != nil { + return fmt.Errorf("scan session transfer authorization: %w", err) + } + if rows.Next() { + return ErrSessionTransferInProgress + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate session transfer authorization: %w", err) + } + if operationID == "" || + operationID != activeOperationID || + fromSessionID != activeFromSessionID || + toSessionID != activeToSessionID { + return ErrSessionTransferInProgress + } + if activeOwnerID != operationID || !leaseExpiresAt.Valid || !leaseExpiresAt.Time.After(time.Now().UTC()) { + return ErrSessionTransferOwnershipLost + } + return nil +} + +func commitAuthorizedSessionTransferTx( + ctx context.Context, + tx *sqlx.Tx, + db *sqlx.DB, + fromSessionID, toSessionID, operationID string, +) error { + if err := authorizeSessionTransferTx( + ctx, tx, db, fromSessionID, toSessionID, operationID, + ); err != nil { + return err + } + return tx.Commit() +} + +// GuardSessionTransferInTransaction rejects mutations that race an active +// durable session transfer. Callers must first hold the queue_session_locks +// row for sessionID. A missing compensation table means no active transfer. +func GuardSessionTransferInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, sessionID string) error { + if err := guardSessionTransferTx(ctx, tx, db, sessionID); err != nil { + if internaldb.IsMissingTableError(err) { + return nil + } + return err + } + return nil +} + +func guardSessionTransferTx(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, sessionID string) error { + var active bool + if err := tx.GetContext(ctx, &active, db.Rebind(` + SELECT EXISTS ( + SELECT 1 FROM queue_session_transfer_compensations + WHERE from_session_id = ? OR to_session_id = ? + ) + `), sessionID, sessionID); err != nil { + return fmt.Errorf("guard active session transfer: %w", err) + } + if active { + return ErrSessionTransferInProgress + } + return nil +} + +// ResolveSessionTransferInTransaction maps a source session to the destination +// selected by the active durable transfer. Callers must first hold the +// queue_session_locks row for sessionID. +func ResolveSessionTransferInTransaction( + ctx context.Context, + tx *sqlx.Tx, + db *sqlx.DB, + sessionID string, +) (string, error) { + rows, err := tx.QueryxContext(ctx, db.Rebind(` + SELECT from_session_id, to_session_id + FROM queue_session_transfer_compensations + WHERE from_session_id = ? OR to_session_id = ? + `), sessionID, sessionID) + if err != nil { + return "", fmt.Errorf("resolve active session transfer: %w", err) + } + defer func() { _ = rows.Close() }() + if !rows.Next() { + if err := rows.Err(); err != nil { + return "", fmt.Errorf("iterate active session transfer: %w", err) + } + return sessionID, nil + } + var fromSessionID, toSessionID string + if err := rows.Scan(&fromSessionID, &toSessionID); err != nil { + return "", fmt.Errorf("scan active session transfer: %w", err) + } + if rows.Next() { + return "", ErrSessionTransferInProgress + } + if err := rows.Err(); err != nil { + return "", fmt.Errorf("iterate active session transfer: %w", err) + } + if sessionID == fromSessionID { + return toSessionID, nil + } + return sessionID, nil +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_sqlite.go b/apps/backend/internal/orchestrator/messagequeue/repository_sqlite.go index f298c1ca5f..f9d9ea668c 100644 --- a/apps/backend/internal/orchestrator/messagequeue/repository_sqlite.go +++ b/apps/backend/internal/orchestrator/messagequeue/repository_sqlite.go @@ -24,12 +24,14 @@ type sqliteRepository struct { mu sync.Mutex sessionLocks map[string]*sync.Mutex - // tasksTablePresent records whether both owning task and session tables - // exist. Production always supplies them; queue-only repository tests use - // the legacy textual methods, while every identity-bound method fails - // closed when session authority is unavailable. Detection happens outside - // transactions because a missing-table statement aborts PostgreSQL txs. - tasksTablePresent bool + // tasksTablePresent is whether the owning tasks table exists, resolved at + // construction. The queue repository's isolated tests create only queue + // tables; production databases always have the task schema. It is checked + // OUTSIDE any transaction because a failed statement on PostgreSQL aborts + // the whole transaction — the guard must never issue its UPDATE against a + // missing table inside a tx. + tasksTablePresent bool + taskSessionsTablePresent bool } // NewSQLiteRepository creates a SQLite-backed Repository. The supplied writer @@ -39,6 +41,12 @@ func NewSQLiteRepository(writer, reader *sqlx.DB) (Repository, error) { if err := r.initSchema(); err != nil { return nil, fmt.Errorf("messagequeue: init schema: %w", err) } + if err := r.ensureSessionTransferCompensationSchema(context.Background()); err != nil { + return nil, fmt.Errorf("messagequeue: init session transfer schema: %w", err) + } + if err := r.ensureEditLeaseSchema(context.Background()); err != nil { + return nil, fmt.Errorf("messagequeue: init edit lease schema: %w", err) + } var present bool var err error if writer.DriverName() == "pgx" { @@ -52,9 +60,27 @@ func NewSQLiteRepository(writer, reader *sqlx.DB) (Repository, error) { return nil, fmt.Errorf("messagequeue: resolve tasks table presence: %w", err) } r.tasksTablePresent = present + if writer.DriverName() == "pgx" { + err = writer.Get(&present, `SELECT to_regclass('task_sessions') IS NOT NULL`) + } else { + err = writer.Get(&present, `SELECT EXISTS (SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'task_sessions')`) + } + if err != nil { + return nil, fmt.Errorf("messagequeue: resolve task_sessions table presence: %w", err) + } + r.taskSessionsTablePresent = present if present { - if err := r.migratePendingMoveIdentities(); err != nil { - return nil, fmt.Errorf("messagequeue: migrate pending move identities: %w", err) + // Older isolated queue fixtures can provide a task_sessions table + // without the incarnation column. They still use the legacy queue API, + // so there is no identity to backfill until the owning schema migrates. + incarnationColumn, columnErr := internaldb.ColumnExists(writer, "task_sessions", "queue_incarnation_id") + if columnErr != nil { + return nil, fmt.Errorf("messagequeue: resolve task session incarnation column: %w", columnErr) + } + if incarnationColumn { + if err := r.migratePendingMoveIdentities(); err != nil { + return nil, fmt.Errorf("messagequeue: migrate pending move identities: %w", err) + } } } return r, nil @@ -99,13 +125,74 @@ func (r *sqliteRepository) migratePendingMoveIdentities() error { // merge-wins/drain-wins ordering deterministic. // // lockSessionTx is the cross-process counterpart: every mutating method takes -// the per-session queue_session_locks row inside its transaction, so tail -// scans and tail changes cannot interleave between backend instances. It is a -// no-op on SQLite (single writer; FOR UPDATE is ignored). +// the per-session queue_session_locks row inside its transaction and rejects +// mutations while a durable transfer owns that session. SQLite serializes +// writers globally; PostgreSQL locks the session row with FOR UPDATE. func (r *sqliteRepository) lockSessionTx(ctx context.Context, tx *sqlx.Tx, sessionID string) error { + if err := r.lockSessionTxUnfenced(ctx, tx, sessionID); err != nil { + return err + } + return guardSessionTransferTx(ctx, tx, r.db, sessionID) +} + +func (r *sqliteRepository) lockSessionTxUnfenced(ctx context.Context, tx *sqlx.Tx, sessionID string) error { return lockSessionTxIn(ctx, tx, r.db, sessionID) } +func (r *sqliteRepository) beginSessionMutationTx( + ctx context.Context, + sessionID, operation string, +) (*sqlx.Tx, error) { + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin %s: %w", operation, err) + } + if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { + _ = tx.Rollback() + return nil, err + } + return tx, nil +} + +func (r *sqliteRepository) withSessionTransferFence( + ctx context.Context, + sessionID string, + fn func(context.Context) error, +) error { + tx, err := r.beginSessionMutationTx(ctx, sessionID, "session transfer fence") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + if err := fn(ctx); err != nil { + return err + } + return tx.Commit() +} + +func (r *sqliteRepository) validateSessionEntryForEditReplay( + ctx context.Context, + sessionID, entryID string, +) error { + tx, err := r.beginSessionMutationTx(ctx, sessionID, "validate edit replay") + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + var present bool + if err := tx.GetContext(ctx, &present, r.db.Rebind(` + SELECT EXISTS ( + SELECT 1 FROM queued_messages WHERE session_id = ? AND id = ? + ) + `), sessionID, entryID); err != nil { + return fmt.Errorf("validate edit replay entry: %w", err) + } + if !present { + return ErrEntryNotFound + } + return tx.Commit() +} + // guardActiveTaskTx rejects queue admissions whose owning task is not live // (archived or deleted). It takes the task-row lock, so admission serializes // with task lifecycle cleanup in the global task-row -> session-lock order and @@ -136,6 +223,85 @@ func (r *sqliteRepository) guardActiveTaskTx(ctx context.Context, tx *sqlx.Tx, t return nil } +// guardSessionTx verifies that the owning task session still exists. The +// session lock acquired before this call already rejects active transfers. +func (r *sqliteRepository) guardSessionTx(ctx context.Context, tx *sqlx.Tx, sessionID, taskID string) error { + return r.guardSessionOwnerTx(ctx, tx, sessionID, taskID) +} + +func (r *sqliteRepository) guardSessionOwnerTx(ctx context.Context, tx *sqlx.Tx, sessionID, taskID string) error { + if !r.tasksTablePresent || !r.taskSessionsTablePresent { + return nil + } + query := `SELECT EXISTS (SELECT 1 FROM task_sessions WHERE id = ?` + args := []interface{}{sessionID} + if taskID != "" { + query += ` AND task_id = ?` + args = append(args, taskID) + } + query += `)` + var present bool + if err := tx.GetContext(ctx, &present, r.db.Rebind(query), args...); err != nil { + return fmt.Errorf("guard task session for queue write: %w", err) + } + if !present { + return ErrTaskInactive + } + return nil +} +func (r *sqliteRepository) captureReservationGenerationsTx( + ctx context.Context, + tx *sqlx.Tx, + msg *QueuedMessage, +) error { + sessionGeneration, err := r.getSendNowGenerationTx(ctx, tx, msg.SessionID) + if err != nil { + return err + } + var lifecycleGeneration int64 + if msg.TaskID != "" { + lifecycleGeneration, err = getLifecycleGenerationTx(ctx, tx, r.db, msg.TaskID) + if err != nil { + return err + } + } + msg.reservationSessionGeneration = sessionGeneration + msg.reservationLifecycleGeneration = lifecycleGeneration + msg.reservationGenerationsCaptured = true + return nil +} + +// LockSessionInTransaction acquires the queue's cross-process session lock +// inside an existing transaction and rejects mutations during an active +// durable transfer. +func LockSessionInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, sessionID string) error { + if err := lockSessionTxIn(ctx, tx, db, sessionID); err != nil { + return err + } + return GuardSessionTransferInTransaction(ctx, tx, db, sessionID) +} + +// LockSessionPairInTransaction acquires both transfer-session locks without +// applying the mutation fence. It is reserved for the owned transfer itself, +// which must mutate attachment claims while its compensation row is active. +func LockSessionPairInTransaction( + ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, firstSessionID, secondSessionID string, +) error { + first, second := firstSessionID, secondSessionID + if first > second { + first, second = second, first + } + if err := lockSessionTxIn(ctx, tx, db, first); err != nil { + return err + } + if first != second { + if err := lockSessionTxIn(ctx, tx, db, second); err != nil { + return err + } + } + return nil +} + func (r *sqliteRepository) validateSessionIdentityTx( ctx context.Context, tx *sqlx.Tx, @@ -150,7 +316,7 @@ func (r *sqliteRepository) validateSessionIdentityTx( FROM task_sessions WHERE id = ? AND task_id = ?` if r.db.DriverName() == "pgx" { - query += ` FOR UPDATE` + query += postgresForUpdateSuffix } err := tx.GetContext( ctx, @@ -310,7 +476,8 @@ func (r *sqliteRepository) initSchema() error { auto_merge_override INTEGER, auto_merge_revision BIGINT NOT NULL DEFAULT 0, send_now_generation BIGINT NOT NULL DEFAULT 0, - status_generation BIGINT NOT NULL DEFAULT 0 + status_generation BIGINT NOT NULL DEFAULT 0, + next_position BIGINT NOT NULL DEFAULT 0 ); `) if err != nil { @@ -340,6 +507,7 @@ func (r *sqliteRepository) initSchema() error { {"Auto-merge revision", `ALTER TABLE queue_session_state ADD COLUMN auto_merge_revision BIGINT NOT NULL DEFAULT 0`}, {"Send Now generation", `ALTER TABLE queue_session_state ADD COLUMN send_now_generation BIGINT NOT NULL DEFAULT 0`}, {"Status generation", `ALTER TABLE queue_session_state ADD COLUMN status_generation BIGINT NOT NULL DEFAULT 0`}, + {"Queue position", `ALTER TABLE queue_session_state ADD COLUMN next_position BIGINT NOT NULL DEFAULT 0`}, } { if _, alterErr := r.db.Exec(migration.sql); alterErr != nil && !internaldb.IsDuplicateColumnError(alterErr) { return fmt.Errorf("add queue state %s: %w", migration.name, alterErr) @@ -347,6 +515,46 @@ func (r *sqliteRepository) initSchema() error { } return nil } +func (r *sqliteRepository) nextQueuePositionTx(ctx context.Context, tx *sqlx.Tx, sessionID string) (int64, error) { + var maxPos sql.NullInt64 + if err := tx.GetContext(ctx, &maxPos, + r.db.Rebind(`SELECT MAX(position) FROM queued_messages WHERE session_id = ?`), sessionID, + ); err != nil { + return 0, fmt.Errorf("max queue position: %w", err) + } + var nextPosition int64 + err := tx.GetContext(ctx, &nextPosition, r.db.Rebind(` + SELECT next_position FROM queue_session_state WHERE session_id = ? + `), sessionID) + if errors.Is(err, sql.ErrNoRows) { + nextPosition = 0 + } else if err != nil { + return 0, fmt.Errorf("read next queue position: %w", err) + } + if maxPos.Valid && maxPos.Int64 > nextPosition { + nextPosition = maxPos.Int64 + } + nextPosition++ + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + INSERT INTO queue_session_state (session_id, next_position) VALUES (?, ?) + ON CONFLICT(session_id) DO UPDATE SET next_position = excluded.next_position + `), sessionID, nextPosition); err != nil { + return 0, fmt.Errorf("advance next queue position: %w", err) + } + return nextPosition, nil +} + +func (r *sqliteRepository) bumpQueuePositionTx(ctx context.Context, tx *sqlx.Tx, sessionID string, position int64) error { + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + INSERT INTO queue_session_state (session_id, next_position) VALUES (?, ?) + ON CONFLICT(session_id) DO UPDATE SET next_position = + CASE WHEN queue_session_state.next_position < excluded.next_position + THEN excluded.next_position ELSE queue_session_state.next_position END + `), sessionID, position); err != nil { + return fmt.Errorf("record queue position: %w", err) + } + return nil +} // Insert appends a new entry at the tail of the session's FIFO queue. func (r *sqliteRepository) Insert(ctx context.Context, msg *QueuedMessage, maxPerSession int) error { @@ -400,6 +608,9 @@ func (r *sqliteRepository) insert( if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil { return err } + if err := r.guardSessionTx(ctx, tx, msg.SessionID, msg.TaskID); err != nil { + return err + } if err := r.validateOptionalSessionIdentityTx(ctx, tx, identity); err != nil { return err } @@ -414,11 +625,11 @@ func (r *sqliteRepository) insert( return err } - var maxPos sql.NullInt64 - if err := tx.GetContext(ctx, &maxPos, r.db.Rebind(`SELECT MAX(position) FROM queued_messages WHERE session_id = ?`), msg.SessionID); err != nil { - return fmt.Errorf("max position: %w", err) + position, err := r.nextQueuePositionTx(ctx, tx, msg.SessionID) + if err != nil { + return err } - msg.Position = maxPos.Int64 + 1 + msg.Position = position if msg.ID == "" { msg.ID = uuid.New().String() } @@ -652,6 +863,9 @@ func (r *sqliteRepository) ensureQueueCapacityTx( // insert. This keeps positions positive for TransferSession and future queue // mutations while preserving the current order. func (r *sqliteRepository) RequeuePreservingFIFO(ctx context.Context, msg *QueuedMessage) error { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return err + } return r.requeuePreservingFIFO(ctx, nil, msg) } @@ -682,6 +896,19 @@ func (r *sqliteRepository) requeuePreservingFIFO( if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil { return err } + if msg.reservationGenerationsCaptured && !msg.IsDurableLifecycle() { + if err := r.validatePendingQueueDispatchTx(ctx, tx, msg); err != nil { + return err + } + } + if err := r.guardSessionTx(ctx, tx, msg.SessionID, msg.TaskID); err != nil { + return err + } + if msg.reservationGenerationsCaptured && !msg.IsDurableLifecycle() { + if err := r.deletePendingQueueDispatchTx(ctx, tx, msg); err != nil { + return err + } + } if identity != nil { if err := r.validateSessionIdentityTx(ctx, tx, *identity); err != nil { return err @@ -810,6 +1037,9 @@ func (r *sqliteRepository) releaseReservedQueueRowTx( // queue position for this coalesce key — supersede→requeue of the same // content keeps FIFO at the same slot. Caller owns the tx. func (r *sqliteRepository) applyCoalesceReplaceTx(ctx context.Context, tx *sqlx.Tx, msg *QueuedMessage, existingID string) error { + if err := r.deleteEditLeaseForEntryTx(ctx, tx, msg.SessionID, existingID); err != nil { + return err + } attachmentsJSON, err := marshalAttachments(msg.Attachments) if err != nil { return err @@ -877,6 +1107,17 @@ func (r *sqliteRepository) applyHeadInsertTx(ctx context.Context, tx *sqlx.Tx, m } else { msg.Position = 1 } + var maxPosition sql.NullInt64 + if err := tx.GetContext(ctx, &maxPosition, r.db.Rebind(`SELECT MAX(position) FROM queued_messages WHERE session_id = ?`), msg.SessionID); err != nil { + return fmt.Errorf("max position after requeue shift: %w", err) + } + if maxPosition.Valid { + if err := r.bumpQueuePositionTx(ctx, tx, msg.SessionID, maxPosition.Int64); err != nil { + return err + } + } else if err := r.bumpQueuePositionTx(ctx, tx, msg.SessionID, msg.Position); err != nil { + return err + } if msg.ID == "" { msg.ID = uuid.New().String() } @@ -907,6 +1148,12 @@ func (r *sqliteRepository) applyHeadInsertTx(ctx context.Context, tx *sqlx.Tx, m // Restore reinserts a previously dequeued entry at its original FIFO position. func (r *sqliteRepository) Restore(ctx context.Context, msg *QueuedMessage, maxPerSession int) error { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return err + } + if err := r.ensureEditLeaseSchema(ctx); err != nil { + return err + } return r.restore(ctx, nil, msg, maxPerSession) } @@ -939,6 +1186,22 @@ func (r *sqliteRepository) restore( if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil { return err } + if msg.reservationGenerationsCaptured && !msg.IsDurableLifecycle() { + if err := r.validatePendingQueueDispatchTx(ctx, tx, msg); err != nil { + return err + } + } + if err := r.guardSessionTx(ctx, tx, msg.SessionID, msg.TaskID); err != nil { + return err + } + if err := r.deleteEditLeaseForEntryTx(ctx, tx, msg.SessionID, msg.ID); err != nil { + return err + } + if msg.reservationGenerationsCaptured && !msg.IsDurableLifecycle() { + if err := r.deletePendingQueueDispatchTx(ctx, tx, msg); err != nil { + return err + } + } if identity != nil { if err := r.validateSessionIdentityTx(ctx, tx, *identity); err != nil { return err @@ -954,6 +1217,9 @@ func (r *sqliteRepository) restore( return ErrQueueFull } } + if err := r.bumpQueuePositionTx(ctx, tx, msg.SessionID, msg.Position); err != nil { + return err + } if msg.ID == "" { msg.ID = uuid.New().String() } @@ -1004,6 +1270,9 @@ func (r *sqliteRepository) appendMatchingTailTx( // AppendOrInsertTail concatenates onto the tail entry when its owner matches, otherwise inserts a new entry. func (r *sqliteRepository) AppendOrInsertTail(ctx context.Context, sessionID, taskID, content, model, queuedBy string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, maxPerSession int) (*QueuedMessage, bool, error) { + if err := r.ensureEditLeaseSchema(ctx); err != nil { + return nil, false, err + } return r.appendOrInsertTail(ctx, nil, sessionID, taskID, content, model, queuedBy, planMode, attachments, metadata, maxPerSession) } @@ -1023,6 +1292,9 @@ func (r *sqliteRepository) appendOrInsertTail(ctx context.Context, identity *Que if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { return nil, false, err } + if err := r.guardSessionTx(ctx, tx, sessionID, taskID); err != nil { + return nil, false, err + } if err := r.validateOptionalSessionIdentityTx(ctx, tx, identity); err != nil { return nil, false, err } @@ -1031,6 +1303,15 @@ func (r *sqliteRepository) appendOrInsertTail(ctx context.Context, identity *Que if err != nil { return nil, false, err } + if tail != nil && tail.QueuedBy == queuedBy { + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, tail.ID) + if err != nil { + return nil, false, err + } + if blocked { + return nil, false, ErrEditConflict + } + } appended, err := r.appendMatchingTailTx(ctx, tx, tail, queuedBy, content) if err != nil { return nil, false, err @@ -1046,16 +1327,16 @@ func (r *sqliteRepository) appendOrInsertTail(ctx context.Context, identity *Que if err := r.ensureQueueCapacityTx(ctx, tx, sessionID, maxPerSession); err != nil { return nil, false, err } - var maxPos sql.NullInt64 - if err := tx.GetContext(ctx, &maxPos, r.db.Rebind(`SELECT MAX(position) FROM queued_messages WHERE session_id = ?`), sessionID); err != nil { - return nil, false, fmt.Errorf("max position: %w", err) + position, err := r.nextQueuePositionTx(ctx, tx, sessionID) + if err != nil { + return nil, false, err } msg := &QueuedMessage{ ID: uuid.New().String(), SessionID: sessionID, TaskID: taskID, - Position: maxPos.Int64 + 1, + Position: position, Content: content, Model: model, PlanMode: planMode, @@ -1090,6 +1371,9 @@ func (r *sqliteRepository) appendOrInsertTail(ctx context.Context, identity *Que // InsertOrReplaceByCoalesceKey replaces an entry with the same session/queued_by/coalesce key, or inserts when allowInsert is set. func (r *sqliteRepository) InsertOrReplaceByCoalesceKey(ctx context.Context, msg *QueuedMessage, coalesceKey string, maxPerSession int, allowInsert bool) (*QueuedMessage, bool, error) { + if err := r.ensureEditLeaseSchema(ctx); err != nil { + return nil, false, err + } return r.insertOrReplaceByCoalesceKey(ctx, nil, msg, coalesceKey, maxPerSession, allowInsert) } @@ -1126,6 +1410,9 @@ func (r *sqliteRepository) insertOrReplaceByCoalesceKey( if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil { return nil, false, err } + if err := r.guardSessionTx(ctx, tx, msg.SessionID, msg.TaskID); err != nil { + return nil, false, err + } if err := r.validateOptionalSessionIdentityTx(ctx, tx, identity); err != nil { return nil, false, err } @@ -1135,6 +1422,9 @@ func (r *sqliteRepository) insertOrReplaceByCoalesceKey( return nil, false, err } if existing != nil { + if err := r.deleteEditLeaseForEntryTx(ctx, tx, msg.SessionID, existing.ID); err != nil { + return nil, false, err + } updated, err := r.replaceCoalesced(ctx, tx, existing, msg) if err != nil { return nil, false, err @@ -1223,6 +1513,9 @@ func (r *sqliteRepository) insertOrReplaceLifecycleByCoalesceKey( if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil { return nil, false, err } + if err := r.guardSessionTx(ctx, tx, msg.SessionID, msg.TaskID); err != nil { + return nil, false, err + } if err := r.validateOptionalSessionIdentityTx(ctx, tx, identity); err != nil { return nil, false, err } @@ -1291,6 +1584,39 @@ func (r *sqliteRepository) LifecycleGeneration(ctx context.Context, taskID strin return generation, nil } +func getLifecycleGenerationTx(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskID string) (int64, error) { + var generation int64 + err := tx.GetContext(ctx, &generation, db.Rebind(` + SELECT generation FROM lifecycle_queue_generations WHERE task_id = ? + `), taskID) + if errors.Is(err, sql.ErrNoRows) { + return 0, nil + } + if err != nil { + return 0, fmt.Errorf("get lifecycle queue generation in transaction: %w", err) + } + return generation, nil +} + +// SessionGeneration returns the current destructive-mutation generation for a +// session. +func (r *sqliteRepository) SessionGeneration(ctx context.Context, sessionID string) (int64, error) { + unlock := r.withSessionLock(sessionID) + defer unlock() + + var generation int64 + err := r.ro.GetContext(ctx, &generation, r.ro.Rebind(` + SELECT send_now_generation FROM queue_session_state WHERE session_id = ? + `), sessionID) + if errors.Is(err, sql.ErrNoRows) { + return 0, nil + } + if err != nil { + return 0, fmt.Errorf("get send-now session generation: %w", err) + } + return generation, nil +} + // PurgeTask removes all task rows and advances its generation. func (r *sqliteRepository) PurgeTask(ctx context.Context, taskID string) (int, error) { tx, err := r.db.BeginTxx(ctx, nil) @@ -1351,11 +1677,54 @@ func purgeQueueRowSessions(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskID return sessions, nil } +func ensureTaskPurgeRecoverySchemas(ctx context.Context, tx *sqlx.Tx) error { + if _, err := tx.ExecContext(ctx, queueDispatchRecoverySchema); err != nil { + return fmt.Errorf("ensure task purge dispatch recovery schema: %w", err) + } + if _, err := tx.ExecContext(ctx, sendNowClaimRecoverySchema); err != nil { + return fmt.Errorf("ensure task purge Send Now recovery schema: %w", err) + } + if _, err := tx.ExecContext(ctx, attachmentCleanupSchema); err != nil { + return fmt.Errorf("ensure task purge attachment cleanup schema: %w", err) + } + if _, err := tx.ExecContext(ctx, sessionTransferCompensationSchema); err != nil { + return fmt.Errorf("ensure task purge transfer compensation schema: %w", err) + } + if _, err := tx.ExecContext(ctx, editLeaseSchema); err != nil { + return fmt.Errorf("ensure task purge edit lease schema: %w", err) + } + return nil +} + +// Attachment cleanup obligations intentionally survive task purge. Archive +// preserves task attachment storage, while already-obsolete queue claims still +// need the retry worker to release them; deletion removes attachment bytes +// through the task attachment service after the database mutation. +func deleteTaskRecoveryRowsTx( + ctx context.Context, + tx *sqlx.Tx, + db *sqlx.DB, + dispatchEntryIDs []string, +) error { + for _, entryID := range dispatchEntryIDs { + if _, err := tx.ExecContext(ctx, db.Rebind(` + DELETE FROM queue_dispatch_claims WHERE entry_id = ? + `), entryID); err != nil { + return fmt.Errorf("purge task dispatch claim: %w", err) + } + } + return nil +} + // PurgeTaskInTransaction lets the task repository make archive/delete and // durable queue invalidation one SQLite transaction. It is backend-internal: // user queue handlers must keep using ownership-checked deletion methods. // // taskSessions is the task's authoritative session set, discovered by the +// caller from its own task_sessions schema. The queue repository never reaches +// across schemas because a failed cross-schema query would abort the caller's +// transaction on PostgreSQL. Standalone purges discover sessions from visible +// queue rows and durable recovery records. // caller from its own task_sessions schema — the queue repository never // reaches across schemas, and a failed cross-schema query would abort the // caller's transaction on PostgreSQL. The standalone PurgeTask passes nil and @@ -1372,7 +1741,7 @@ func DeleteSessionInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, i var incarnationID string query := `SELECT queue_incarnation_id FROM task_sessions WHERE id = ? AND task_id = ?` if db.DriverName() == "pgx" { - query += ` FOR UPDATE` + query += postgresForUpdateSuffix } err := tx.GetContext(ctx, &incarnationID, db.Rebind(query), identity.SessionID, identity.TaskID) if errors.Is(err, sql.ErrNoRows) { @@ -1384,14 +1753,8 @@ func DeleteSessionInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, i if incarnationID != identity.SessionIncarnationID { return ErrSessionIdentityMismatch } - for _, statement := range []string{ - `DELETE FROM queued_messages WHERE session_id = ?`, - `DELETE FROM pending_moves WHERE session_id = ?`, - `DELETE FROM queue_session_state WHERE session_id = ?`, - } { - if _, err := tx.ExecContext(ctx, db.Rebind(statement), identity.SessionID); err != nil { - return fmt.Errorf("delete queue session state: %w", err) - } + if _, err := PurgeSessionInTransaction(ctx, tx, db, identity.SessionID); err != nil { + return fmt.Errorf("delete queue session state: %w", err) } return nil } @@ -1414,6 +1777,9 @@ func PurgeTaskInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskI } } + if err := ensureTaskPurgeRecoverySchemas(ctx, tx); err != nil { + return 0, err + } // Serialize with per-session tail operations: a purge that races a fold // or insert could otherwise delete a row an admission just accepted, or // admit into a queue being purged. Lock the AUTHORITATIVE session set — @@ -1428,6 +1794,16 @@ func PurgeTaskInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskI if err != nil { return 0, err } + dispatchEntryIDs, dispatchSessions, err := pendingQueueDispatchesForTaskTx(ctx, tx, taskID) + if err != nil { + return 0, err + } + sendNowSessions, err := pendingSendNowSessionsForTaskTx(ctx, tx, taskID) + if err != nil { + return 0, err + } + rowSessions = append(rowSessions, sendNowSessions...) + rowSessions = append(rowSessions, dispatchSessions...) seen := make(map[string]struct{}, len(rowSessions)+len(taskSessions)) for _, sessionID := range rowSessions { seen[sessionID] = struct{}{} @@ -1445,6 +1821,14 @@ func PurgeTaskInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskI return 0, err } } + for _, sessionID := range ordered { + if err := guardSessionTransferTx(ctx, tx, db, sessionID); err != nil { + return 0, err + } + } + if err := deleteEditLeasesForSessionIDsTx(ctx, tx, db, ordered); err != nil { + return 0, err + } result, err := tx.ExecContext(ctx, db.Rebind(`DELETE FROM queued_messages WHERE task_id = ?`), taskID) if err != nil { return 0, fmt.Errorf("purge queued task entries: %w", err) @@ -1453,13 +1837,19 @@ func PurgeTaskInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskI if err != nil { return 0, fmt.Errorf("purge queued task entries rows affected: %w", err) } - if _, err := tx.ExecContext( - ctx, - db.Rebind(`DELETE FROM pending_moves WHERE task_id = ?`), - taskID, - ); err != nil { + if _, err := tx.ExecContext(ctx, db.Rebind(`DELETE FROM pending_moves WHERE task_id = ?`), taskID); err != nil { return 0, fmt.Errorf("purge pending task moves: %w", err) } + if err := deleteTaskRecoveryRowsTx(ctx, tx, db, dispatchEntryIDs); err != nil { + return 0, err + } + for _, sessionID := range ordered { + if _, err := tx.ExecContext(ctx, db.Rebind(` + DELETE FROM queue_send_now_claims WHERE session_id = ? + `), sessionID); err != nil { + return 0, fmt.Errorf("purge task Send Now claim: %w", err) + } + } if _, err := tx.ExecContext(ctx, db.Rebind(` INSERT INTO lifecycle_queue_generations (task_id, generation) VALUES (?, 1) ON CONFLICT(task_id) DO UPDATE SET generation = lifecycle_queue_generations.generation + 1 @@ -1469,6 +1859,71 @@ func PurgeTaskInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskI return int(removed), nil } +// PurgeSessionInTransaction removes all durable queue state for a deleted +// session while the owning task repository transaction is still open. The +// caller must hold the session lock and must commit the surrounding +// transaction after this function returns successfully. +func PurgeSessionInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, sessionID string) (int, error) { + for _, table := range []string{"queued_messages", "queue_session_locks"} { + present, err := internaldb.TableExists(tx, table) + if err != nil { + return 0, fmt.Errorf("check %s table: %w", table, err) + } + if !present { + return 0, nil + } + } + + if err := ensureTaskPurgeRecoverySchemas(ctx, tx); err != nil { + return 0, err + } + if err := lockSessionTxIn(ctx, tx, db, sessionID); err != nil { + return 0, err + } + if err := guardSessionTransferTx(ctx, tx, db, sessionID); err != nil { + return 0, err + } + if err := deleteEditLeasesForSessionIDsTx(ctx, tx, db, []string{sessionID}); err != nil { + return 0, err + } + + result, err := tx.ExecContext(ctx, db.Rebind(` + DELETE FROM queued_messages WHERE session_id = ? + `), sessionID) + if err != nil { + return 0, fmt.Errorf("purge queued session entries: %w", err) + } + removed, err := result.RowsAffected() + if err != nil { + return 0, fmt.Errorf("purge queued session entries rows affected: %w", err) + } + if _, err := tx.ExecContext(ctx, db.Rebind(` + DELETE FROM pending_moves WHERE session_id = ? + `), sessionID); err != nil { + return 0, fmt.Errorf("purge session pending move: %w", err) + } + if _, err := tx.ExecContext(ctx, db.Rebind(` + INSERT INTO queue_session_state (session_id, auto_run, send_now_generation) + VALUES (?, 1, 1) + ON CONFLICT(session_id) DO UPDATE SET + auto_run = 1, + send_now_generation = queue_session_state.send_now_generation + 1 + `), sessionID); err != nil { + return 0, fmt.Errorf("advance deleted session queue generation: %w", err) + } + if _, err := tx.ExecContext(ctx, db.Rebind(` + DELETE FROM queue_send_now_claims WHERE session_id = ? + `), sessionID); err != nil { + return 0, fmt.Errorf("purge session Send Now claim: %w", err) + } + if _, err := tx.ExecContext(ctx, db.Rebind(` + DELETE FROM queue_dispatch_claims WHERE session_id = ? + `), sessionID); err != nil { + return 0, fmt.Errorf("purge session dispatch claims: %w", err) + } + return int(removed), nil +} + // replaceCoalesced overwrites the existing coalesced row with msg inside the transaction. func (r *sqliteRepository) replaceCoalesced(ctx context.Context, tx *sqlx.Tx, existing, msg *QueuedMessage) (*QueuedMessage, error) { if msg.QueuedAt.IsZero() { @@ -1517,11 +1972,11 @@ func (r *sqliteRepository) insertCoalesced(ctx context.Context, tx *sqlx.Tx, msg if err := r.ensureQueueCapacity(ctx, tx, msg.SessionID, maxPerSession); err != nil { return err } - var maxPos sql.NullInt64 - if err := tx.GetContext(ctx, &maxPos, r.db.Rebind(`SELECT MAX(position) FROM queued_messages WHERE session_id = ?`), msg.SessionID); err != nil { - return fmt.Errorf("max position: %w", err) + position, err := r.nextQueuePositionTx(ctx, tx, msg.SessionID) + if err != nil { + return err } - msg.Position = maxPos.Int64 + 1 + msg.Position = position if msg.ID == "" { msg.ID = uuid.New().String() } @@ -1783,6 +2238,23 @@ func (r *sqliteRepository) ListDurableLifecycleEntries(ctx context.Context) ([]Q return out, nil } +func (r *sqliteRepository) FindByID(ctx context.Context, entryID string) (*QueuedMessage, error) { + row := r.ro.QueryRowxContext(ctx, r.ro.Rebind(` + SELECT id, session_id, task_id, position, content, model, plan_mode, + attachments_json, metadata_json, queued_at, queued_by + FROM queued_messages + WHERE id = ? + `), entryID) + message, err := scanQueuedRow(row) + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrEntryNotFound + } + if err != nil { + return nil, fmt.Errorf("find queued entry: %w", err) + } + return message, nil +} + // CountBySession returns the number of entries for a session. func (r *sqliteRepository) CountBySession(ctx context.Context, sessionID string) (int, error) { var n int @@ -1839,6 +2311,9 @@ func (r *sqliteRepository) TakeHead(ctx context.Context, sessionID string) (*Que // the same queue are serialized in-process, not just at the DB layer. unlock := r.withSessionLock(sessionID) defer unlock() + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return nil, err + } tx, err := r.db.BeginTxx(ctx, nil) if err != nil { @@ -1864,12 +2339,27 @@ func (r *sqliteRepository) TakeHead(ctx context.Context, sessionID string) (*Que } return nil, fmt.Errorf("take head: %w", err) } + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, msg.ID) + if err != nil { + return nil, err + } + if blocked { + return nil, nil + } + if err := r.captureReservationGenerationsTx(ctx, tx, msg); err != nil { + return nil, err + } // Two concurrent TakeHead calls can both observe the same head row in // their respective DEFERRED transactions; one wins the DELETE, the other // finds RowsAffected()==0 after waiting on the writer lock. Returning the // already-drained message in that case would let the orchestrator dispatch // it twice. Treat the lost race as "queue empty for now" so the caller // retries on the next agent.ready instead. + if !msg.IsDurableLifecycle() { + if err := r.persistQueueDispatchClaimTx(ctx, tx, msg); err != nil { + return nil, err + } + } res, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM queued_messages WHERE id = ?`), msg.ID) if err != nil { return nil, fmt.Errorf("delete head: %w", err) @@ -2147,6 +2637,31 @@ func (r *sqliteRepository) setAutoRunTx( } return nil } +func (r *sqliteRepository) getSendNowGenerationTx(ctx context.Context, tx *sqlx.Tx, sessionID string) (int64, error) { + var generation int64 + err := tx.GetContext(ctx, &generation, r.db.Rebind(` + SELECT send_now_generation FROM queue_session_state WHERE session_id = ? + `), sessionID) + if errors.Is(err, sql.ErrNoRows) { + return 0, nil + } + if err != nil { + return 0, fmt.Errorf("get send-now generation: %w", err) + } + return generation, nil +} + +func (r *sqliteRepository) bumpSendNowGenerationTx(ctx context.Context, tx *sqlx.Tx, sessionID string) error { + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + INSERT INTO queue_session_state (session_id, send_now_generation) + VALUES (?, 1) + ON CONFLICT(session_id) DO UPDATE + SET send_now_generation = queue_session_state.send_now_generation + 1 + `), sessionID); err != nil { + return fmt.Errorf("advance send-now generation: %w", err) + } + return nil +} func (r *sqliteRepository) getAutoRunTx( ctx context.Context, @@ -2224,6 +2739,9 @@ func (r *sqliteRepository) reserveHead( ) (*QueuedMessage, bool, error) { unlock := r.withSessionLock(sessionID) defer unlock() + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return nil, true, err + } tx, err := r.db.BeginTxx(ctx, nil) if err != nil { @@ -2296,6 +2814,16 @@ func (r *sqliteRepository) reserveHeadTx( continue } retainedLifecycle := msg.IsDurableLifecycle() + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, msg.ID) + if err != nil { + return nil, true, err + } + if blocked { + return nil, true, nil + } + if err := r.captureReservationGenerationsTx(ctx, tx, msg); err != nil { + return nil, true, err + } if retainedLifecycle || retainOrdinary { reserved, err := r.reserveRetainedHead( ctx, tx, identity, msg, storedMetadataJSON, retainedLifecycle, @@ -2350,10 +2878,11 @@ func (r *sqliteRepository) reserveRetainedHead( // Strip metadata persisted by an interrupted prior process from the // returned copy so a failed retry becomes visible again. msg.Metadata = clearReservedMetadata(msg.Metadata) - reservedMetadata := markReservedMetadata(msg.Metadata) + msg.lifecycleReservationID = uuid.NewString() + reservedMetadata := markReservedMetadata(msg.Metadata, msg.lifecycleReservationID) if identity != nil { reservedMetadata = markReservedMetadataForIncarnation( - msg.Metadata, + reservedMetadata, identity.SessionIncarnationID, ) } @@ -2390,6 +2919,9 @@ func (r *sqliteRepository) reserveOrdinaryHead( tx *sqlx.Tx, msg *QueuedMessage, ) (*QueuedMessage, error) { + if err := r.persistQueueDispatchClaimTx(ctx, tx, msg); err != nil { + return nil, err + } res, err := tx.ExecContext(ctx, r.db.Rebind(` DELETE FROM queued_messages WHERE id = ? AND session_id = ? `), msg.ID, msg.SessionID) @@ -2409,39 +2941,99 @@ func (r *sqliteRepository) reserveOrdinaryHead( return msg, nil } -// AcknowledgeByID removes a reserved durable entry after executor acceptance. -func (r *sqliteRepository) AcknowledgeByID(ctx context.Context, sessionID, entryID string) error { - unlock := r.withSessionLock(sessionID) +// AcknowledgeReserved removes only the exact lifecycle reservation accepted by +// the executor. A newer retry with the same queue entry ID must survive. +func (r *sqliteRepository) AcknowledgeReserved(ctx context.Context, msg *QueuedMessage) error { + unlock := r.withSessionLock(msg.SessionID) defer unlock() - // The ack must serialize with restore/replace on the same session across - // processes: an autocommit DELETE could otherwise delete a row that a - // concurrent backend just restored/reinserted, losing durable retry state. tx, err := r.db.BeginTxx(ctx, nil) if err != nil { return fmt.Errorf("begin acknowledge tx: %w", err) } defer func() { _ = tx.Rollback() }() - if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { + if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil { return err } - res, err := tx.ExecContext(ctx, r.db.Rebind(` - DELETE FROM queued_messages WHERE id = ? AND session_id = ? - `), entryID, sessionID) + var storedMetadataJSON string + err = tx.GetContext(ctx, &storedMetadataJSON, r.db.Rebind(` + SELECT metadata_json FROM queued_messages WHERE id = ? AND session_id = ? + `), msg.ID, msg.SessionID) + if errors.Is(err, sql.ErrNoRows) { + return ErrEntryNotFound + } + if err != nil { + return fmt.Errorf("read reserved lifecycle entry: %w", err) + } + storedMetadata := make(map[string]interface{}) + if storedMetadataJSON != "" && storedMetadataJSON != "{}" { + if err := json.Unmarshal([]byte(storedMetadataJSON), &storedMetadata); err != nil { + return fmt.Errorf("unmarshal reserved lifecycle metadata: %w", err) + } + } + storedReservationID, _ := storedMetadata[metadataLifecycleReservationID].(string) + if msg.lifecycleReservationID == "" || storedReservationID != msg.lifecycleReservationID { + return ErrLifecycleReservationChanged + } + result, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queued_messages + WHERE id = ? AND session_id = ? AND metadata_json = ? + `), msg.ID, msg.SessionID, storedMetadataJSON) if err != nil { return fmt.Errorf("acknowledge queued: %w", err) } - affected, err := res.RowsAffected() + affected, err := result.RowsAffected() if err != nil { return err } - if affected == 0 { - return ErrEntryNotFound + if affected != 1 { + return ErrLifecycleReservationChanged } - if err := tx.Commit(); err != nil { + return tx.Commit() +} + +func (r *sqliteRepository) AcknowledgeByID(ctx context.Context, sessionID, entryID string) error { + unlock := r.withSessionLock(sessionID) + defer unlock() + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin acknowledge-by-id tx: %w", err) + } + defer func() { _ = tx.Rollback() }() + if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { return err } - return nil + var metadataJSON string + err = tx.GetContext(ctx, &metadataJSON, r.db.Rebind(` + SELECT metadata_json FROM queued_messages WHERE id = ? AND session_id = ? + `), entryID, sessionID) + if errors.Is(err, sql.ErrNoRows) { + return ErrEntryNotFound + } + if err != nil { + return fmt.Errorf("read acknowledge-by-id entry: %w", err) + } + reserved, err := isReservedMetadataJSON(metadataJSON) + if err != nil { + return err + } + if !reserved { + return ErrEntryNotFound + } + result, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queued_messages WHERE id = ? AND session_id = ? AND metadata_json = ? + `), entryID, sessionID, metadataJSON) + if err != nil { + return fmt.Errorf("acknowledge-by-id: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected != 1 { + return ErrEntryNotFound + } + return tx.Commit() } func (r *sqliteRepository) AcknowledgeByIDForSession( @@ -2637,6 +3229,9 @@ func (r *sqliteRepository) takeByID( ) (*QueuedMessage, error) { unlock := r.withSessionLock(sessionID) defer unlock() + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return nil, err + } tx, err := r.db.BeginTxx(ctx, nil) if err != nil { @@ -2663,6 +3258,21 @@ func (r *sqliteRepository) takeByID( } return nil, fmt.Errorf("take by id: %w", err) } + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, entryID) + if err != nil { + return nil, err + } + if blocked { + return nil, ErrEditConflict + } + if err := r.captureReservationGenerationsTx(ctx, tx, msg); err != nil { + return nil, err + } + if !msg.IsDurableLifecycle() { + if err := r.persistQueueDispatchClaimTx(ctx, tx, msg); err != nil { + return nil, err + } + } res, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM queued_messages WHERE id = ? AND session_id = ?`), msg.ID, sessionID) if err != nil { return nil, fmt.Errorf("delete by id: %w", err) @@ -2720,6 +3330,9 @@ func (r *sqliteRepository) claimSendNow(ctx context.Context, identity *QueueSess if len(expected) == 0 { return nil, ErrSendNowEmpty } + if err := r.ensureSendNowClaimRecoverySchema(ctx); err != nil { + return nil, err + } unlock := r.withSessionLock(sessionID) defer unlock() @@ -2751,16 +3364,37 @@ func (r *sqliteRepository) claimSendNow(ctx context.Context, identity *QueueSess if err != nil { return nil, err } + for _, source := range sources { + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, source.ID) + if err != nil { + return nil, err + } + if blocked { + return nil, ErrEditConflict + } + } claimIdentity := queueSessionIdentityValue(identity) bindSendNowLifecycleReservations(sources, claimIdentity) envelope, err := BuildSendNowEnvelope(sources) if err != nil { return nil, err } + sessionGeneration, err := r.getSendNowGenerationTx(ctx, tx, sessionID) + if err != nil { + return nil, err + } generations, err := lifecycleGenerationsForSourcesTx(ctx, tx, r.db, sources) if err != nil { return nil, err } + claim := &SendNowClaim{ + ClaimID: uuid.NewString(), + Identity: claimIdentity, + Sources: sources, + Dispatch: *envelope, + SourceGenerations: generations, + SessionGeneration: sessionGeneration, + } if err := r.applySQLiteSendNowClaim(ctx, tx, identity, sessionID, sources, storedByID); err != nil { return nil, err @@ -2772,21 +3406,22 @@ func (r *sqliteRepository) claimSendNow(ctx context.Context, identity *QueueSess if err != nil { return nil, err } + claim.OperationGeneration = operationGeneration + claim.SessionGeneration = operationGeneration + if err := r.persistSendNowClaimTx(ctx, tx, claim); err != nil { + return nil, err + } if err := tx.Commit(); err != nil { return nil, err } - return &SendNowClaim{ - Identity: claimIdentity, - OperationGeneration: operationGeneration, - Sources: sources, - Dispatch: *envelope, - SourceGenerations: generations, - }, nil + return claim, nil } +const sendNowRestoreAction = "restore" + // RestoreSendNowClaim puts every claimed source back at its original position. func (r *sqliteRepository) RestoreSendNowClaim(ctx context.Context, claim *SendNowClaim) error { - tx, sessionID, unlock, err := r.beginSendNowClaimTx(ctx, claim, "restore") + tx, sessionID, unlock, err := r.beginSendNowClaimTx(ctx, claim, sendNowRestoreAction) if err != nil { return err } @@ -2796,6 +3431,11 @@ func (r *sqliteRepository) RestoreSendNowClaim(ctx context.Context, claim *SendN if err != nil { return err } + sessionGeneration, err := r.getSendNowGenerationTx(ctx, tx, sessionID) + if err != nil { + return err + } + sessionChanged := claim.SessionGeneration != sessionGeneration generations := make(map[string]int64) for _, source := range claim.Sources { if source.TaskID == "" { @@ -2807,6 +3447,9 @@ func (r *sqliteRepository) RestoreSendNowClaim(ctx context.Context, claim *SendN } generations[source.TaskID] = generation } + if sessionChanged && !sendNowClaimSourcesAllInvalidated(claim, generations) { + return ErrSendNowClaimChanged + } if err := validateSQLiteSendNowRestore(claim, sessionID, stored, generations); err != nil { return err } @@ -2819,11 +3462,19 @@ func (r *sqliteRepository) RestoreSendNowClaim(ctx context.Context, claim *SendN return err } } + if claim.ClaimID != "" { + if err := r.deleteSendNowClaimTx(ctx, tx, sessionID); err != nil { + return err + } + } return tx.Commit() } // AcknowledgeSendNowClaim removes every durable source after the replacement prompt is accepted. func (r *sqliteRepository) AcknowledgeSendNowClaim(ctx context.Context, claim *SendNowClaim) error { + if claim != nil && claim.ClaimID == "" { + return ErrSendNowClaimChanged + } tx, sessionID, unlock, err := r.beginSendNowClaimTx(ctx, claim, "acknowledge") if err != nil { return err @@ -2834,18 +3485,55 @@ func (r *sqliteRepository) AcknowledgeSendNowClaim(ctx context.Context, claim *S if err != nil { return err } - if err := validateSQLiteSendNowAcknowledge(claim, sessionID, stored); err != nil { + sessionGeneration, err := r.getSendNowGenerationTx(ctx, tx, sessionID) + if err != nil { + return err + } + sessionChanged := claim.SessionGeneration != sessionGeneration + generations := make(map[string]int64) + for _, source := range claim.Sources { + if source.TaskID == "" { + continue + } + generation, generationErr := lifecycleGenerationInTx(ctx, tx, r.db, source.TaskID) + if generationErr != nil { + return generationErr + } + generations[source.TaskID] = generation + } + if sessionChanged && !sendNowClaimSourcesAllInvalidated(claim, generations) { + if claim.ClaimID != "" { + if err := r.deleteExactSendNowClaimTx(ctx, tx, sessionID, claim.ClaimID); err != nil { + return err + } + } + if err := tx.Commit(); err != nil { + return err + } + return ErrSendNowClaimChanged + } + if err := validateSQLiteSendNowAcknowledge(claim, sessionID, stored, generations); err != nil { return err } for _, source := range claim.Sources { + if sendNowSourceGenerationChanged(claim, source, generations[source.TaskID]) { + continue + } if err := r.acknowledgeSQLiteSendNowSource(ctx, tx, sessionID, source, stored); err != nil { return err } } + if claim.ClaimID != "" { + if err := r.deleteExactSendNowClaimTx(ctx, tx, sessionID, claim.ClaimID); err != nil { + return err + } + } return tx.Commit() } // beginSendNowClaimTx starts the claim transaction and takes the per-session lock. +// +//nolint:nestif // Claim fencing must keep the transaction and unlock paths paired. func (r *sqliteRepository) beginSendNowClaimTx( ctx context.Context, claim *SendNowClaim, @@ -2870,6 +3558,44 @@ func (r *sqliteRepository) beginSendNowClaimTx( unlock() return nil, "", nil, err } + if claim.ClaimID != "" { + var currentClaimID string + err := tx.QueryRowxContext(ctx, r.db.Rebind(` + SELECT claim_id FROM queue_send_now_claims WHERE session_id = ? + `), sessionID).Scan(¤tClaimID) + if errors.Is(err, sql.ErrNoRows) && action == sendNowRestoreAction { + var movedSessionID string + movedErr := tx.QueryRowxContext(ctx, r.db.Rebind(` + SELECT session_id FROM queue_send_now_claims WHERE claim_id = ? + `), claim.ClaimID).Scan(&movedSessionID) + if movedErr == nil || !errors.Is(movedErr, sql.ErrNoRows) { + _ = tx.Rollback() + unlock() + if movedErr != nil { + return nil, "", nil, fmt.Errorf("validate transferred send-now %s claim: %w", action, movedErr) + } + return nil, "", nil, ErrSendNowClaimChanged + } + } + if errors.Is(err, sql.ErrNoRows) && action != sendNowRestoreAction { + _ = tx.Rollback() + unlock() + return nil, "", nil, ErrSendNowClaimChanged + } + if err == nil && currentClaimID != claim.ClaimID { + _ = tx.Rollback() + unlock() + return nil, "", nil, ErrSendNowClaimChanged + } + if err != nil && !errors.Is(err, sql.ErrNoRows) { + _ = tx.Rollback() + unlock() + return nil, "", nil, fmt.Errorf("validate send-now %s claim: %w", action, err) + } + } + // A task purge deletes the durable claim row along with its source rows. + // Restore validates the session and task generations below, so it can + // distinguish that destructive invalidation from an unrelated claim race. if claim.Identity.SessionIncarnationID != "" { if err := r.validateSessionIdentityTx(ctx, tx, claim.Identity); err != nil { _ = tx.Rollback() @@ -2890,6 +3616,19 @@ func (r *sqliteRepository) beginSendNowClaimTx( return nil, "", nil, fmt.Errorf("validate send-now generation: %w", err) } if errors.Is(err, sql.ErrNoRows) || operationGeneration != claim.OperationGeneration { + if err == nil { + if deleteErr := r.deleteExactSendNowClaimTx(ctx, tx, sessionID, claim.ClaimID); deleteErr != nil { + _ = tx.Rollback() + unlock() + return nil, "", nil, deleteErr + } + if commitErr := tx.Commit(); commitErr != nil { + unlock() + return nil, "", nil, commitErr + } + unlock() + return nil, "", nil, ErrSendNowClaimChanged + } _ = tx.Rollback() unlock() return nil, "", nil, ErrSendNowClaimChanged @@ -3034,9 +3773,9 @@ func (r *sqliteRepository) reserveSQLiteSendNowSource( source QueuedMessage, stored storedQueueEntry, ) error { - metadata := markReservedMetadata(source.Metadata) + metadata := markReservedMetadata(source.Metadata, uuid.NewString()) if identity != nil { - metadata = markReservedMetadataForIncarnation(source.Metadata, identity.SessionIncarnationID) + metadata = markReservedMetadataForIncarnation(metadata, identity.SessionIncarnationID) } metadataJSON, err := marshalMetadata(metadata) if err != nil { @@ -3154,6 +3893,9 @@ func (r *sqliteRepository) insertSQLiteSendNowSource(ctx context.Context, tx *sq if err != nil { return err } + if err := r.bumpQueuePositionTx(ctx, tx, source.SessionID, source.Position); err != nil { + return err + } metadataJSON, err := marshalMetadata(clearReservedMetadata(source.Metadata)) if err != nil { return err @@ -3170,15 +3912,24 @@ func (r *sqliteRepository) insertSQLiteSendNowSource(ctx context.Context, tx *sq } // validateSQLiteSendNowAcknowledge verifies every durable source is still reserved before acknowledgement. -func validateSQLiteSendNowAcknowledge(claim *SendNowClaim, sessionID string, stored map[string]storedQueueEntry) error { +func validateSQLiteSendNowAcknowledge( + claim *SendNowClaim, + sessionID string, + stored map[string]storedQueueEntry, + generations map[string]int64, +) error { for _, source := range claim.Sources { if source.SessionID != sessionID { return ErrSendNowClaimChanged } + if sendNowSourceGenerationChanged(claim, source, generations[source.TaskID]) { + continue + } if !source.IsDurableLifecycle() { continue } - if entry, ok := stored[source.ID]; ok && !entry.message.IsReservedInFlight() { + entry, ok := stored[source.ID] + if !ok || !entry.message.IsReservedInFlight() { return ErrSendNowClaimChanged } } @@ -3257,7 +4008,7 @@ func (r *sqliteRepository) UpdateContent(ctx context.Context, sessionID, entryID // UpdateContentAndMetadata replaces content and applies metadata updates to an entry owned by queuedBy. func (r *sqliteRepository) UpdateContentAndMetadata(ctx context.Context, sessionID, entryID, content string, attachments []MessageAttachment, metadataUpdates map[string]interface{}, queuedBy string) error { - return r.updateContentAndMetadata(ctx, nil, sessionID, entryID, content, attachments, metadataUpdates, queuedBy, nil) + return r.updateContentAndMetadataWithLease(ctx, sessionID, entryID, "", content, attachments, metadataUpdates, queuedBy) } func (r *sqliteRepository) UpdateContentAndMetadataForSession(ctx context.Context, identity QueueSessionIdentity, entryID, content string, attachments []MessageAttachment, metadataUpdates map[string]interface{}, queuedBy string) error { @@ -3294,7 +4045,6 @@ func (r *sqliteRepository) updateContentAndMetadata(ctx context.Context, identit if err := r.validateOptionalSessionIdentityTx(ctx, tx, identity); err != nil { return err } - var metadataJSON string query := `SELECT metadata_json FROM queued_messages WHERE id = ? AND session_id = ? AND queued_by = ? AND queued_by NOT IN (?, ?, ?)` if err := tx.GetContext(ctx, &metadataJSON, r.db.Rebind(query), entryID, sessionID, queuedBy, QueuedByAgent, QueuedByWorkflow, QueuedByServer); err != nil { @@ -3317,14 +4067,110 @@ func (r *sqliteRepository) updateContentAndMetadata(ctx context.Context, identit if err != nil { return err } - if err := r.updateQueuedEntryTx( - ctx, tx, sessionID, entryID, content, attachmentsJSON, metadataJSON, queuedBy, - ); err != nil { + if err := r.updateQueuedEntryTx(ctx, tx, sessionID, entryID, content, attachmentsJSON, metadataJSON, queuedBy); err != nil { return err } return tx.Commit() } +func (r *sqliteRepository) authorizeQueueEditLeaseTx( + ctx context.Context, + tx *sqlx.Tx, + sessionID, entryID, leaseID string, +) error { + if leaseID == "" { + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, entryID) + if err != nil { + return err + } + if blocked { + return ErrEditConflict + } + return nil + } + var authorized bool + if err := tx.GetContext(ctx, &authorized, tx.Rebind(` + SELECT EXISTS ( + SELECT 1 FROM queue_edit_leases + WHERE session_id = ? AND entry_id = ? AND lease_id = ? AND expires_at > ? + ) + `), sessionID, entryID, leaseID, time.Now().UTC()); err != nil { + return fmt.Errorf("authorize queue edit lease: %w", err) + } + if !authorized { + return ErrEditLeaseNotFound + } + return nil +} + +func (r *sqliteRepository) updateContentAndMetadataTx( + ctx context.Context, + tx *sqlx.Tx, + sessionID, entryID, content, queuedBy string, + attachmentsJSON string, + metadataUpdates map[string]interface{}, +) error { + var metadataJSON string + query := `SELECT metadata_json FROM queued_messages WHERE id = ? AND session_id = ? AND queued_by = ? AND queued_by NOT IN (?, ?, ?)` + if err := tx.GetContext(ctx, &metadataJSON, r.db.Rebind(query), entryID, sessionID, queuedBy, QueuedByAgent, QueuedByWorkflow, QueuedByServer); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrEntryNotFound + } + return fmt.Errorf("read queued metadata: %w", err) + } + var metadata map[string]interface{} + if metadataJSON != "" && metadataJSON != "{}" { + if err := json.Unmarshal([]byte(metadataJSON), &metadata); err != nil { + return fmt.Errorf("unmarshal queued metadata: %w", err) + } + } + metadataJSON, err := marshalMetadata(applyMetadataUpdates(metadata, metadataUpdates)) + if err != nil { + return err + } + res, err := tx.ExecContext(ctx, r.db.Rebind(`UPDATE queued_messages SET content = ?, attachments_json = ?, metadata_json = ? + WHERE id = ? AND session_id = ? AND queued_by = ? AND queued_by NOT IN (?, ?, ?)`), + content, attachmentsJSON, metadataJSON, entryID, sessionID, queuedBy, QueuedByAgent, QueuedByWorkflow, QueuedByServer) + if err != nil { + return fmt.Errorf("update queued: %w", err) + } + n, err := res.RowsAffected() + if err != nil { + return err + } + if n == 0 { + return ErrEntryNotFound + } + return tx.Commit() +} + +func (r *sqliteRepository) updateContentAndMetadataWithLease(ctx context.Context, sessionID, entryID, leaseID, content string, attachments []MessageAttachment, metadataUpdates map[string]interface{}, queuedBy string) error { + if queuedBy == "" || IsReservedQueuedBy(queuedBy) { + return ErrEntryNotFound + } + if err := r.ensureEditLeaseSchema(ctx); err != nil { + return err + } + attachmentsJSON, err := marshalAttachments(attachments) + if err != nil { + return err + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin update queued tx: %w", err) + } + defer func() { _ = tx.Rollback() }() + if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { + return err + } + if err := r.authorizeQueueEditLeaseTx(ctx, tx, sessionID, entryID, leaseID); err != nil { + return err + } + return r.updateContentAndMetadataTx( + ctx, tx, sessionID, entryID, content, queuedBy, attachmentsJSON, metadataUpdates, + ) +} + // MergeIntoAbove folds the source entry into the entry directly above it within // the same session, updating the target row and deleting the source row in one // transaction so a concurrent drain can never observe a half-merged queue. @@ -3372,6 +4218,15 @@ func (r *sqliteRepository) mergeIntoAbove(ctx context.Context, identity *QueueSe if !mergeAllowed(source, target, queuedBy) { return nil, ErrNoMergeTarget } + for _, entryID := range []string{source.ID, target.ID} { + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, entryID) + if err != nil { + return nil, err + } + if blocked { + return nil, ErrEditConflict + } + } content, attachments, metadata, err := buildMergedEntry(target, source) if err != nil { @@ -3396,6 +4251,9 @@ func (r *sqliteRepository) mergeIntoAbove(ctx context.Context, identity *QueueSe // predecessor in one transaction. Missing or incompatible candidates are // successful skips and leave storage unchanged. func (r *sqliteRepository) AutoMergeIntoAbove(ctx context.Context, sessionID, sourceID string) (*QueuedMessage, bool, error) { + if err := r.ensureEditLeaseSchema(ctx); err != nil { + return nil, false, err + } return r.autoMergeIntoAbove(ctx, nil, sessionID, sourceID, nil) } @@ -3457,6 +4315,15 @@ func (r *sqliteRepository) autoMergeIntoAbove( if !compatible { return source, false, nil } + for _, entryID := range []string{source.ID, target.ID} { + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, entryID) + if err != nil { + return nil, false, err + } + if blocked { + return nil, false, ErrEditConflict + } + } if err := applyMergeWrites(ctx, r, tx, target, source, values.content, values.attachments, values.metadata, sessionID); err != nil { if errors.Is(err, ErrEntryNotFound) { return nil, false, nil @@ -3479,6 +4346,9 @@ func (r *sqliteRepository) autoMergeIntoAbove( // the fold is the admission. Missing or incompatible tails are successful // skips and leave storage unchanged. func (r *sqliteRepository) AutoMergeCandidateIntoAbove(ctx context.Context, candidate *QueuedMessage) (*QueuedMessage, bool, error) { + if err := r.ensureEditLeaseSchema(ctx); err != nil { + return nil, false, err + } return r.autoMergeCandidateIntoAbove(ctx, nil, candidate, nil, nil) } @@ -3536,6 +4406,9 @@ func (r *sqliteRepository) autoMergeCandidateIntoAbove( if err := r.lockSessionTx(ctx, tx, candidate.SessionID); err != nil { return nil, false, err } + if err := r.guardSessionTx(ctx, tx, candidate.SessionID, candidate.TaskID); err != nil { + return nil, false, err + } if err := r.validateOptionalSessionIdentityTx(ctx, tx, identity); err != nil { return nil, false, err } @@ -3560,6 +4433,13 @@ func (r *sqliteRepository) autoMergeCandidateIntoAbove( if !compatible { return nil, false, nil } + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, candidate.SessionID, target.ID) + if err != nil { + return nil, false, err + } + if blocked { + return nil, false, ErrEditConflict + } if err := claimOptionalMessageAttachmentsTx( ctx, tx, identity, claim, candidate.TaskID, candidate.SessionID, ); err != nil { @@ -3698,6 +4578,15 @@ func (r *sqliteRepository) reorderEntries(ctx context.Context, identity *QueueSe if err != nil { return err } + for _, msg := range visible { + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, msg.ID) + if err != nil { + return err + } + if blocked { + return ErrEditConflict + } + } // Interleave reserved rows at their current places; visible rows emit in // the submitted order. Positions are compacted to 1..N. @@ -3892,6 +4781,13 @@ func (r *sqliteRepository) deleteByID(ctx context.Context, identity *QueueSessio if reserved { return nil, ErrEntryNotFound } + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, entryID) + if err != nil { + return nil, err + } + if blocked { + return nil, ErrEditConflict + } res, err := tx.ExecContext(ctx, r.db.Rebind(` DELETE FROM queued_messages @@ -3967,6 +4863,9 @@ func (r *sqliteRepository) DeleteAllBySessionForIdentity(ctx context.Context, id func (r *sqliteRepository) deleteAllBySession(ctx context.Context, identity *QueueSessionIdentity, sessionID string) (*QueueRemovalResult, error) { unlock := r.withSessionLock(sessionID) defer unlock() + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return nil, err + } tx, err := r.db.BeginTxx(ctx, nil) if err != nil { @@ -3991,6 +4890,15 @@ func (r *sqliteRepository) deleteAllBySession(ctx context.Context, identity *Que if err != nil { return nil, err } + for _, candidate := range candidates { + blocked, err := r.editLeaseBlocksEntryTx(ctx, tx, sessionID, candidate.message.ID) + if err != nil { + return nil, err + } + if blocked { + return nil, ErrEditConflict + } + } result := &QueueRemovalResult{} for _, candidate := range candidates { if candidate.reserved { @@ -4014,12 +4922,78 @@ func (r *sqliteRepository) deleteAllBySession(ctx context.Context, identity *Que } result.Removed = append(result.Removed, candidate.message) } + if err := r.bumpSendNowGenerationTx(ctx, tx, sessionID); err != nil { + return nil, err + } + if err := r.deletePendingQueueDispatchesBySessionTx(ctx, tx, sessionID); err != nil { + return nil, err + } if err := tx.Commit(); err != nil { return nil, err } return result, nil } +// PurgeSession removes all queue rows for a deleted session, including +// reserved lifecycle deliveries, and its pending workflow move. +func (r *sqliteRepository) PurgeSession(ctx context.Context, sessionID string) (int, error) { + unlock := r.withSessionLock(sessionID) + defer unlock() + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return 0, err + } + if err := r.ensureSendNowClaimRecoverySchema(ctx); err != nil { + return 0, err + } + + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return 0, fmt.Errorf("begin purge session queue tx: %w", err) + } + defer func() { _ = tx.Rollback() }() + if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { + return 0, err + } + if err := r.deleteEditLeasesForSessionTx(ctx, tx, sessionID); err != nil { + return 0, err + } + res, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queued_messages WHERE session_id = ? + `), sessionID) + if err != nil { + return 0, fmt.Errorf("purge session queued messages: %w", err) + } + removed, err := res.RowsAffected() + if err != nil { + return 0, fmt.Errorf("purge session queue rows affected: %w", err) + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM pending_moves WHERE session_id = ? + `), sessionID); err != nil { + return 0, fmt.Errorf("purge session pending move: %w", err) + } + // A deleted session must not retain an explicit OFF policy if its queue + // state is later observed during cleanup or an ID is reused. Keep the + // generation row for send-now fencing, but reset the independent policy to + // its default ON value before advancing that generation. + if err := r.setAutoRunTx(ctx, tx, nil, sessionID, true); err != nil { + return 0, err + } + if err := r.bumpSendNowGenerationTx(ctx, tx, sessionID); err != nil { + return 0, err + } + if err := r.deleteSendNowClaimTx(ctx, tx, sessionID); err != nil { + return 0, err + } + if err := r.deletePendingQueueDispatchesBySessionTx(ctx, tx, sessionID); err != nil { + return 0, err + } + if err := tx.Commit(); err != nil { + return 0, err + } + return int(removed), nil +} + type cancellationCandidate struct { message QueuedMessage metadataJSON string @@ -4087,6 +5061,44 @@ func lifecycleReservationFromMetadataJSON(metadataJSON string) (string, bool, er return lifecycleReservationIncarnation(metadata), message.IsReservedInFlight(), nil } +func (r *sqliteRepository) transferSessionOwned( + ctx context.Context, + oldSessionID, newSessionID, operationID string, +) error { + return r.transferSessionOwnedTx(ctx, oldSessionID, newSessionID, operationID) +} + +func (r *sqliteRepository) beginAuthorizedSessionTransferTx( + ctx context.Context, + oldSessionID, newSessionID, operationID string, +) (*sqlx.Tx, error) { + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin transfer tx: %w", err) + } + if _, _, err := r.lockSessionTransferPairTx(ctx, tx, oldSessionID, newSessionID); err != nil { + _ = tx.Rollback() + return nil, err + } + if err := authorizeSessionTransferTx( + ctx, tx, r.db, oldSessionID, newSessionID, operationID, + ); err != nil { + _ = tx.Rollback() + return nil, err + } + if err := r.guardSessionOwnerTx(ctx, tx, oldSessionID, ""); err != nil { + _ = tx.Rollback() + return nil, err + } + if oldSessionID != newSessionID { + if err := r.guardSessionOwnerTx(ctx, tx, newSessionID, ""); err != nil { + _ = tx.Rollback() + return nil, err + } + } + return tx, nil +} + func (r *sqliteRepository) guardTransferTasksTx( ctx context.Context, tx *sqlx.Tx, @@ -4144,6 +5156,89 @@ func (r *sqliteRepository) lockTransferSessions(oldSessionID, newSessionID strin } } +func (r *sqliteRepository) transferSessionRecoveryRowsTx( + ctx context.Context, + tx *sqlx.Tx, + oldSessionID, newSessionID string, +) error { + if err := r.transferPendingQueueDispatchesTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + UPDATE queue_attachment_cleanups + SET current_session_id = ? + WHERE current_session_id = ? + `), newSessionID, oldSessionID); err != nil { + return fmt.Errorf("transfer attachment cleanup session: %w", err) + } + return nil +} + +func (r *sqliteRepository) transferSessionQueueRowsTx( + ctx context.Context, + tx *sqlx.Tx, + oldSessionID, newSessionID string, +) error { + var destinationMax sql.NullInt64 + if err := tx.GetContext(ctx, &destinationMax, r.db.Rebind(` + SELECT MAX(position) FROM queued_messages WHERE session_id = ? + `), newSessionID); err != nil { + return fmt.Errorf("transfer max: %w", err) + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + UPDATE queued_messages + SET session_id = ?, position = position + ? + WHERE session_id = ? + `), newSessionID, destinationMax.Int64, oldSessionID); err != nil { + return fmt.Errorf("transfer queued: %w", err) + } + var transferredMax sql.NullInt64 + if err := tx.GetContext(ctx, &transferredMax, r.db.Rebind(` + SELECT MAX(position) FROM queued_messages WHERE session_id = ? + `), newSessionID); err != nil { + return fmt.Errorf("transfer destination max after move: %w", err) + } + if transferredMax.Valid { + return r.bumpQueuePositionTx(ctx, tx, newSessionID, transferredMax.Int64) + } + return nil +} + +func (r *sqliteRepository) transferSessionStateTx( + ctx context.Context, + tx *sqlx.Tx, + oldSessionID, newSessionID string, +) error { + sourceAutoRun, err := r.getAutoRunTx(ctx, tx, nil, oldSessionID) + if err != nil { + return err + } + destinationAutoRun, err := r.getAutoRunTx(ctx, tx, nil, newSessionID) + if err != nil { + return err + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM pending_moves WHERE session_id = ? + `), newSessionID); err != nil { + return fmt.Errorf("clear dest pending move: %w", err) + } + if err := r.transferPendingMoveTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + if err := r.setAutoRunTx(ctx, tx, nil, newSessionID, sourceAutoRun && destinationAutoRun); err != nil { + return err + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(` + DELETE FROM queue_session_state WHERE session_id = ? + `), oldSessionID); err != nil { + return fmt.Errorf("clear source queue auto-run: %w", err) + } + if err := r.bumpSendNowGenerationTx(ctx, tx, oldSessionID); err != nil { + return err + } + return r.bumpSendNowGenerationTx(ctx, tx, newSessionID) +} + // TransferSession moves all entries (and any pending move) from one session to another. func (r *sqliteRepository) TransferSession(ctx context.Context, oldSessionID, newSessionID string) error { return r.transferSession(ctx, nil, nil, oldSessionID, newSessionID) @@ -4153,21 +5248,27 @@ func (r *sqliteRepository) TransferSessionIdentities(ctx context.Context, source return r.transferSession(ctx, &source, &destination, source.SessionID, destination.SessionID) } -//nolint:cyclop // transfer preserves queue, policy, and attachment invariants atomically. -func (r *sqliteRepository) transferSession(ctx context.Context, source, destination *QueueSessionIdentity, oldSessionID, newSessionID string) error { - // The transfer moves rows out of the source and into the destination, so - // it must hold BOTH sessions' locks: a concurrent source-side insert - // (holding the source lock) could otherwise commit a row the transfer's - // READ COMMITTED UPDATE missed, orphaning it on the old session. Acquire - // in stable sorted order — in-process and cross-process alike — so - // concurrent transfers in opposite directions cannot deadlock. +//nolint:cyclop,funlen // transfer preserves queue, policy, recovery, and attachment invariants atomically. +func (r *sqliteRepository) transferSession( + ctx context.Context, + source, destination *QueueSessionIdentity, + oldSessionID, newSessionID string, +) error { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return err + } + if err := r.ensureSendNowClaimRecoverySchema(ctx); err != nil { + return err + } + if err := r.ensureAttachmentCleanupSchema(ctx); err != nil { + return err + } first, second, release := r.lockTransferSessions(oldSessionID, newSessionID) defer release() attachmentsTablePresent, err := r.sharedTablePresent("task_message_attachments") if err != nil { return err } - tx, err := r.db.BeginTxx(ctx, nil) if err != nil { return fmt.Errorf("begin transfer tx: %w", err) @@ -4198,24 +5299,16 @@ func (r *sqliteRepository) transferSession(ctx context.Context, source, destinat return err } } - - // Shift positions on the destination so transferred entries land at the tail - // without colliding on the (session_id, position) implicit ordering. - var destMax sql.NullInt64 - if err := tx.GetContext(ctx, &destMax, r.db.Rebind(`SELECT MAX(position) FROM queued_messages WHERE session_id = ?`), newSessionID); err != nil { - return fmt.Errorf("transfer max: %w", err) + if err := r.transferSessionRecoveryRowsTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err } - if _, err := tx.ExecContext(ctx, r.db.Rebind(` - UPDATE queued_messages - SET session_id = ?, position = position + ? - WHERE session_id = ? - `), newSessionID, destMax.Int64, oldSessionID); err != nil { - return fmt.Errorf("transfer queued: %w", err) + if err := r.deleteEditLeasesForTransferTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err } - - if _, err := tx.ExecContext(ctx, r.db.Rebind(` - DELETE FROM pending_moves WHERE session_id = ? - `), newSessionID); err != nil { + if err := r.transferSessionQueueRowsTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + if _, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM pending_moves WHERE session_id = ?`), newSessionID); err != nil { return fmt.Errorf("clear dest pending move: %w", err) } if err := r.transferPendingMoveTx(ctx, tx, oldSessionID, newSessionID); err != nil { @@ -4224,14 +5317,101 @@ func (r *sqliteRepository) transferSession(ctx context.Context, source, destinat if err := r.setAutoRunTx(ctx, tx, destination, newSessionID, sourceAutoRun && destinationAutoRun); err != nil { return err } - if _, err := tx.ExecContext(ctx, r.db.Rebind(` - DELETE FROM queue_session_state WHERE session_id = ? - `), oldSessionID); err != nil { + if _, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM queue_session_state WHERE session_id = ?`), oldSessionID); err != nil { return fmt.Errorf("clear source queue auto-run: %w", err) } + if err := r.bumpSendNowGenerationTx(ctx, tx, oldSessionID); err != nil { + return err + } + if err := r.bumpSendNowGenerationTx(ctx, tx, newSessionID); err != nil { + return err + } + if err := r.transferPendingSendNowClaimTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } return tx.Commit() } +func (r *sqliteRepository) transferSessionOwnedTx( + ctx context.Context, + oldSessionID, newSessionID, operationID string, +) error { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return err + } + if err := r.ensureSendNowClaimRecoverySchema(ctx); err != nil { + return err + } + if err := r.ensureAttachmentCleanupSchema(ctx); err != nil { + return err + } + // The transfer moves rows out of the source and into the destination, so + // it must hold BOTH sessions' locks: a concurrent source-side insert + // (holding the source lock) could otherwise commit a row the transfer's + // READ COMMITTED UPDATE missed, orphaning it on the old session. Acquire + // in stable sorted order — in-process and cross-process alike — so + // concurrent transfers in opposite directions cannot deadlock. + first, second := oldSessionID, newSessionID + if first > second { + first, second = second, first + } + unlockFirst := r.withSessionLock(first) + defer unlockFirst() + if first != second { + unlockSecond := r.withSessionLock(second) + defer unlockSecond() + } + + tx, err := r.beginAuthorizedSessionTransferTx( + ctx, oldSessionID, newSessionID, operationID, + ) + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + if oldSessionID == newSessionID { + return commitAuthorizedSessionTransferTx( + ctx, tx, r.db, oldSessionID, newSessionID, operationID, + ) + } + if err := r.transferSessionRecoveryRowsTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + if err := r.deleteEditLeasesForTransferTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + if err := r.transferSessionQueueRowsTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + if err := r.transferSessionStateTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + if err := r.transferPendingSendNowClaimTx(ctx, tx, oldSessionID, newSessionID); err != nil { + return err + } + return commitAuthorizedSessionTransferTx( + ctx, tx, r.db, oldSessionID, newSessionID, operationID, + ) +} + +func queuedSnapshotTaskIDs(entries []QueuedMessage, pendingMove *PendingMove) []string { + taskIDs := make(map[string]struct{}, len(entries)+1) + for _, entry := range entries { + if entry.TaskID != "" { + taskIDs[entry.TaskID] = struct{}{} + } + } + if pendingMove != nil && pendingMove.TaskID != "" { + taskIDs[pendingMove.TaskID] = struct{}{} + } + ids := make([]string, 0, len(taskIDs)) + for taskID := range taskIDs { + ids = append(ids, taskID) + } + sort.Strings(ids) + return ids +} + //nolint:cyclop // every malformed or conflicting attachment reference must roll back transfer. func (r *sqliteRepository) transferQueuedAttachmentClaimsTx( ctx context.Context, @@ -4314,16 +5494,32 @@ func rejectReservedLifecycleTransferTx( tx *sqlx.Tx, sessionID string, ) error { + claimedIDs := make(map[string]struct{}) + var claimJSON string + err := tx.QueryRowxContext(ctx, tx.Rebind(` + SELECT claim_json FROM queue_send_now_claims WHERE session_id = ? + `), sessionID).Scan(&claimJSON) + if err == nil { + var claim SendNowClaim + if unmarshalErr := json.Unmarshal([]byte(claimJSON), &claim); unmarshalErr != nil { + return fmt.Errorf("unmarshal Send Now claim before transfer: %w", unmarshalErr) + } + for _, source := range claim.Sources { + claimedIDs[source.ID] = struct{}{} + } + } else if !errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("read Send Now claim before transfer: %w", err) + } rows, err := tx.QueryxContext(ctx, tx.Rebind(` - SELECT metadata_json FROM queued_messages WHERE session_id = ? + SELECT id, metadata_json FROM queued_messages WHERE session_id = ? `), sessionID) if err != nil { return fmt.Errorf("list lifecycle reservations before transfer: %w", err) } defer func() { _ = rows.Close() }() for rows.Next() { - var metadataJSON string - if err := rows.Scan(&metadataJSON); err != nil { + var entryID, metadataJSON string + if err := rows.Scan(&entryID, &metadataJSON); err != nil { return fmt.Errorf("scan lifecycle reservation before transfer: %w", err) } reserved, err := isReservedMetadataJSON(metadataJSON) @@ -4331,6 +5527,9 @@ func rejectReservedLifecycleTransferTx( return err } if reserved { + if _, claimed := claimedIDs[entryID]; claimed { + continue + } return ErrQueueChanged } } @@ -4353,7 +5552,6 @@ func (r *sqliteRepository) transferPendingMoveTx( } return nil } - var destinationIncarnationID string if err := tx.GetContext(ctx, &destinationIncarnationID, r.db.Rebind(` SELECT queue_incarnation_id FROM task_sessions WHERE id = ? @@ -4466,37 +5664,69 @@ func (r *sqliteRepository) restorePendingMoveTx( } func (r *sqliteRepository) replaceSession(ctx context.Context, identity *QueueSessionIdentity, sessionID string, entries []QueuedMessage, pendingMove *PendingMove) error { + if err := r.ensureQueueDispatchRecoverySchema(ctx); err != nil { + return err + } + if err := r.ensureEditLeaseSchema(ctx); err != nil { + return err + } tx, err := r.db.BeginTxx(ctx, nil) if err != nil { return fmt.Errorf("begin replace session tx: %w", err) } defer func() { _ = tx.Rollback() }() + if err := validateReplacementSnapshot(identity, entries, pendingMove); err != nil { + return err + } + for _, taskID := range queuedSnapshotTaskIDs(entries, pendingMove) { + if err := r.guardActiveTaskTx(ctx, tx, taskID); err != nil { + return err + } + } if err := r.guardOptionalActiveTaskTx(ctx, tx, identity); err != nil { return err } if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { return err } + if err := r.guardSessionTx(ctx, tx, sessionID, ""); err != nil { + return err + } if err := r.validateOptionalSessionIdentityTx(ctx, tx, identity); err != nil { return err } - if err := validateReplacementSnapshot(identity, entries, pendingMove); err != nil { + if err := r.deleteEditLeasesForSessionTx(ctx, tx, sessionID); err != nil { + return err + } + if err := r.deletePendingQueueDispatchesBySessionTx(ctx, tx, sessionID); err != nil { return err } if _, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM queued_messages WHERE session_id = ?`), sessionID); err != nil { return fmt.Errorf("clear queued messages: %w", err) } + var maxPosition int64 + for _, entry := range entries { + if entry.Position > maxPosition { + maxPosition = entry.Position + } + } if err := r.restoreQueueEntriesTx(ctx, tx, sessionID, entries); err != nil { return err } + if err := r.bumpQueuePositionTx(ctx, tx, sessionID, maxPosition); err != nil { + return err + } if _, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM pending_moves WHERE session_id = ?`), sessionID); err != nil { return fmt.Errorf("clear pending move: %w", err) } if err := r.restorePendingMoveTx(ctx, tx, sessionID, pendingMove); err != nil { return err } + if err := r.bumpSendNowGenerationTx(ctx, tx, sessionID); err != nil { + return err + } return tx.Commit() } @@ -4505,16 +5735,8 @@ func (r *sqliteRepository) SetPendingMove(ctx context.Context, sessionID string, if move == nil || move.TaskID == "" || sessionID == "" { return ErrSessionIdentityMismatch } - if move.SessionIncarnationID == "" { - if r.tasksTablePresent { - identity, err := r.ResolveSessionIdentity(ctx, move.TaskID, sessionID) - if err != nil { - return err - } - move.SessionIncarnationID = identity.SessionIncarnationID - } else { - move.SessionIncarnationID = "legacy:" + sessionID - } + if move.SessionIncarnationID == "" && !r.tasksTablePresent { + move.SessionIncarnationID = "legacy:" + sessionID } if move.QueuedAt.IsZero() { move.QueuedAt = time.Now().UTC() @@ -4534,7 +5756,22 @@ func (r *sqliteRepository) SetPendingMove(ctx context.Context, sessionID string, if err := r.lockSessionTx(ctx, tx, sessionID); err != nil { return err } + if err := r.guardSessionTx(ctx, tx, sessionID, ""); err != nil { + return err + } if r.tasksTablePresent { + if move.SessionIncarnationID == "" { + var incarnationID string + if err := tx.GetContext(ctx, &incarnationID, r.db.Rebind(` + SELECT queue_incarnation_id FROM task_sessions WHERE id = ? AND task_id = ? + `), sessionID, move.TaskID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrTaskInactive + } + return fmt.Errorf("resolve pending move session identity: %w", err) + } + move.SessionIncarnationID = incarnationID + } if err := r.validateSessionIdentityTx(ctx, tx, QueueSessionIdentity{ TaskID: move.TaskID, SessionID: sessionID, diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_sqlite_lifecycle_regression_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_sqlite_lifecycle_regression_test.go new file mode 100644 index 0000000000..0da6ee6ef4 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_sqlite_lifecycle_regression_test.go @@ -0,0 +1,39 @@ +package messagequeue + +import ( + "context" + "database/sql" + "testing" + + "github.com/jmoiron/sqlx" + _ "github.com/mattn/go-sqlite3" + "github.com/stretchr/testify/require" +) + +func TestSQLiteReplaceSessionRejectsArchivedSnapshotEntries(t *testing.T) { + raw, err := sql.Open("sqlite3", "file:queue-lifecycle-regression?mode=memory&cache=shared&_foreign_keys=on") + require.NoError(t, err) + raw.SetMaxOpenConns(1) + raw.SetMaxIdleConns(1) + db := sqlx.NewDb(raw, "sqlite3") + t.Cleanup(func() { _ = db.Close() }) + + _, err = db.Exec(` + CREATE TABLE tasks (id TEXT PRIMARY KEY, archived_at DATETIME, updated_at DATETIME); + CREATE TABLE task_sessions (id TEXT PRIMARY KEY, task_id TEXT NOT NULL); + INSERT INTO tasks (id, archived_at, updated_at) VALUES ('task-archived', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP); + INSERT INTO task_sessions (id, task_id) VALUES ('session-archived', 'task-archived'); + `) + require.NoError(t, err) + + repo, err := NewSQLiteRepository(db, db) + require.NoError(t, err) + err = repo.ReplaceSession(context.Background(), "session-archived", []QueuedMessage{ + {ID: "restored", SessionID: "session-archived", TaskID: "task-archived", Position: 1, Content: "stale"}, + }, nil) + require.ErrorIs(t, err, ErrTaskInactive) + err = repo.ReplaceSession(context.Background(), "session-archived", nil, &PendingMove{ + TaskID: "task-archived", + }) + require.ErrorIs(t, err, ErrTaskInactive) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_sqlite_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_sqlite_test.go index 1f585b36fe..36fa74df71 100644 --- a/apps/backend/internal/orchestrator/messagequeue/repository_sqlite_test.go +++ b/apps/backend/internal/orchestrator/messagequeue/repository_sqlite_test.go @@ -692,7 +692,7 @@ func TestSQLiteRepository_DeletePreservesReservedLifecycleEntry(t *testing.T) { if removed != 0 { t.Fatalf("clear removed %d reserved entries, want 0", removed) } - if err := repo.AcknowledgeByID(ctx, "s1", msg.ID); err != nil { + if err := repo.AcknowledgeReserved(ctx, reserved); err != nil { t.Fatalf("acknowledge reserved entry after cancellation attempts: %v", err) } } @@ -715,7 +715,7 @@ func TestSQLiteRepository_CancellationReservationOrdering(t *testing.T) { t.Run("reservation wins before clear", func(t *testing.T) { repo := newTestSQLiteRepo(t) ctx := context.Background() - msg := insertDurableLifecycleEntry(t, repo, "reserve-first") + _ = insertDurableLifecycleEntry(t, repo, "reserve-first") reserved, err := repo.ReserveHead(ctx, "reserve-first") if err != nil || reserved == nil { @@ -725,11 +725,56 @@ func TestSQLiteRepository_CancellationReservationOrdering(t *testing.T) { if err != nil || removed != 0 { t.Fatalf("clear after reserve: removed=%d err=%v", removed, err) } - if err := repo.AcknowledgeByID(ctx, "reserve-first", msg.ID); err != nil { + if err := repo.AcknowledgeReserved(ctx, reserved); err != nil { t.Fatalf("acknowledge: %v", err) } }) } +func TestSQLiteRepository_PurgeSessionRemovesReservedLifecycleAndPendingMove(t *testing.T) { + repo := newTestSQLiteRepo(t) + sqlRepo := repo.(*sqliteRepository) + ctx := context.Background() + _ = insertDurableLifecycleEntry(t, repo, "purge-session") + if err := repo.SetPendingMove(ctx, "purge-session", &PendingMove{TaskID: "t1"}); err != nil { + t.Fatalf("set pending move: %v", err) + } + if err := repo.SetAutoRun(ctx, "purge-session", false); err != nil { + t.Fatalf("set auto-run: %v", err) + } + if reserved, err := repo.ReserveHead(ctx, "purge-session"); err != nil || reserved == nil { + t.Fatalf("reserve lifecycle entry: msg=%+v err=%v", reserved, err) + } + + removed, err := repo.PurgeSession(ctx, "purge-session") + if err != nil { + t.Fatalf("purge session: %v", err) + } + if removed != 1 { + t.Fatalf("purged rows = %d, want 1", removed) + } + entries, err := repo.ListBySession(ctx, "purge-session") + if err != nil { + t.Fatalf("list purged session: %v", err) + } + if len(entries) != 0 { + t.Fatalf("purged session retained entries: %+v", entries) + } + move, err := repo.GetPendingMove(ctx, "purge-session") + if err != nil { + t.Fatalf("get purged pending move: %v", err) + } + if move != nil { + t.Fatalf("purged session retained pending move: %+v", move) + } + var autoRun int + if err := sqlRepo.db.GetContext(ctx, &autoRun, + `SELECT auto_run FROM queue_session_state WHERE session_id = ?`, "purge-session"); err != nil { + t.Fatalf("read purged session policy: %v", err) + } + if autoRun != 1 { + t.Fatalf("purged session retained auto-run=%d, want 1", autoRun) + } +} func insertDurableLifecycleEntry(t *testing.T, repo Repository, sessionID string) *QueuedMessage { t.Helper() diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_task_purge_recovery_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_task_purge_recovery_test.go new file mode 100644 index 0000000000..fee956990f --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_task_purge_recovery_test.go @@ -0,0 +1,281 @@ +package messagequeue + +import ( + "context" + "testing" +) + +func TestSQLiteTaskPurgeRemovesOrdinaryDispatchClaim(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + source := insertTestEntry(t, repo, "session-purge-dispatch", "task-purge-dispatch", "prompt", QueuedByUser, nil, nil) + reserved, _, err := repo.ReserveHeadIfAutoRun(ctx, source.SessionID) + if err != nil || reserved == nil { + t.Fatalf("reserve ordinary head = %#v, err=%v", reserved, err) + } + + tx, err := repo.db.BeginTxx(ctx, nil) + if err != nil { + t.Fatal(err) + } + removed, err := PurgeTaskInTransaction(ctx, tx, repo.db, source.TaskID, []string{source.SessionID}) + if err != nil { + _ = tx.Rollback() + t.Fatal(err) + } + if removed != 0 { + t.Fatalf("removed visible rows = %d, want 0", removed) + } + if err := tx.Commit(); err != nil { + t.Fatal(err) + } + pending, err := repo.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("dispatch claims after task purge = %#v, want empty", pending) + } +} + +func TestSQLiteTaskPurgePreservesAttachmentCleanup(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + cleanup := AttachmentCleanup{ + SessionID: "session-purge-cleanup", EntryID: "entry-purge-cleanup", + OperationID: "operation-purge-cleanup", TaskID: "task-purge-cleanup", + Attachments: []MessageAttachment{{AttachmentID: "attachment-purge-cleanup"}}, + } + if err := repo.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + + tx, err := repo.db.BeginTxx(ctx, nil) + if err != nil { + t.Fatal(err) + } + if _, err := PurgeTaskInTransaction(ctx, tx, repo.db, cleanup.TaskID, []string{cleanup.SessionID}); err != nil { + _ = tx.Rollback() + t.Fatal(err) + } + if err := tx.Commit(); err != nil { + t.Fatal(err) + } + pending, err := repo.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || pending[0].OperationID != cleanup.OperationID { + t.Fatalf("attachment cleanups after task purge = %#v, want retained obligation", pending) + } +} + +func TestSQLiteTaskPurgePreservesOtherTaskRecoveryInSharedSession(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + const ( + sessionID = "session-shared-purge" + taskA = "task-shared-purge-a" + taskB = "task-shared-purge-b" + ) + entryA := insertTestEntry(t, repo, sessionID, taskA, "prompt-a", QueuedByUser, nil, nil) + entryB := insertTestEntry(t, repo, sessionID, taskB, "prompt-b", QueuedByUser, nil, nil) + for _, wantID := range []string{entryA.ID, entryB.ID} { + reserved, _, err := repo.ReserveHeadIfAutoRun(ctx, sessionID) + if err != nil || reserved == nil || reserved.ID != wantID { + t.Fatalf("reserve %s = %#v, err=%v", wantID, reserved, err) + } + } + for _, cleanup := range []AttachmentCleanup{ + { + SessionID: sessionID, EntryID: entryA.ID, OperationID: "cleanup-a", TaskID: taskA, + Attachments: []MessageAttachment{{AttachmentID: "attachment-a"}}, + }, + { + SessionID: sessionID, EntryID: entryB.ID, OperationID: "cleanup-b", TaskID: taskB, + Attachments: []MessageAttachment{{AttachmentID: "attachment-b"}}, + }, + } { + if err := repo.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + } + + tx, err := repo.db.BeginTxx(ctx, nil) + if err != nil { + t.Fatal(err) + } + if _, err := PurgeTaskInTransaction(ctx, tx, repo.db, taskA, []string{sessionID}); err != nil { + _ = tx.Rollback() + t.Fatal(err) + } + if err := tx.Commit(); err != nil { + t.Fatal(err) + } + + claims, err := repo.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(claims) != 1 || claims[0].Message.ID != entryB.ID { + t.Fatalf("remaining dispatch claims = %#v, want only %s", claims, entryB.ID) + } + cleanups, err := repo.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(cleanups) != 2 { + t.Fatalf("remaining attachment cleanups = %#v, want both durable obligations", cleanups) + } +} + +func TestSQLiteTaskPurgeCreatesAbsentRecoveryTables(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + tx, err := repo.db.BeginTxx(ctx, nil) + if err != nil { + t.Fatal(err) + } + if _, err := PurgeTaskInTransaction( + ctx, tx, repo.db, "task-without-recovery-tables", []string{"session-without-recovery-tables"}, + ); err != nil { + _ = tx.Rollback() + t.Fatal(err) + } + if err := tx.Commit(); err != nil { + t.Fatal(err) + } + var claims int + if err := repo.db.GetContext(ctx, &claims, `SELECT COUNT(*) FROM queue_dispatch_claims`); err != nil { + t.Fatalf("query recovery table after purge: %v", err) + } +} + +func TestSQLiteRepositoryTaskPurgeDiscoversOrdinaryDispatchSession(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + source := insertTestEntry(t, repo, "session-direct-purge", "task-direct-purge", "prompt", QueuedByUser, nil, nil) + reserved, _, err := repo.ReserveHeadIfAutoRun(ctx, source.SessionID) + if err != nil || reserved == nil { + t.Fatalf("reserve ordinary head = %#v, err=%v", reserved, err) + } + if _, err := repo.PurgeTask(ctx, source.TaskID); err != nil { + t.Fatal(err) + } + pending, err := repo.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("dispatch claims after direct task purge = %#v, want empty", pending) + } +} + +func TestSQLiteAttachmentCleanupMigratesAndTransfersCurrentSession(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + _, err := repo.db.ExecContext(ctx, ` + CREATE TABLE queue_attachment_cleanups ( + session_id TEXT NOT NULL, + entry_id TEXT NOT NULL, + operation_id TEXT NOT NULL DEFAULT '', + task_id TEXT NOT NULL, + owner_id TEXT NOT NULL DEFAULT '', + lease_id TEXT NOT NULL DEFAULT '', + remove_entry INTEGER NOT NULL DEFAULT 0, + claim_pending INTEGER NOT NULL DEFAULT 0, + entry_fingerprint TEXT NOT NULL DEFAULT '', + attachments_json TEXT NOT NULL DEFAULT '[]', + created_at TIMESTAMP NOT NULL, + PRIMARY KEY (session_id, entry_id, operation_id) + ) + `) + if err != nil { + t.Fatal(err) + } + const ( + sourceSessionID = "legacy-cleanup-source" + destinationSessionID = "legacy-cleanup-destination" + ) + if _, err := repo.db.ExecContext(ctx, ` + INSERT INTO queue_attachment_cleanups + (session_id, entry_id, operation_id, task_id, attachments_json, created_at) + VALUES (?, 'legacy-entry', 'legacy-operation', 'legacy-task', '[]', CURRENT_TIMESTAMP) + `, sourceSessionID); err != nil { + t.Fatal(err) + } + + cleanups, err := repo.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(cleanups) != 1 || cleanups[0].CurrentSessionID != sourceSessionID { + t.Fatalf("migrated cleanup = %#v", cleanups) + } + if err := repo.TransferSession(ctx, sourceSessionID, destinationSessionID); err != nil { + t.Fatal(err) + } + cleanups, err = repo.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(cleanups) != 1 || + cleanups[0].SessionID != sourceSessionID || + cleanups[0].CurrentSessionID != destinationSessionID { + t.Fatalf("transferred cleanup = %#v", cleanups) + } +} + +func TestSQLiteStaleAttachmentCleanupUpsertPreservesTransferredSession(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + cleanup := AttachmentCleanup{ + SessionID: "cleanup-stale-upsert-old", EntryID: "cleanup-stale-upsert-entry", + OperationID: "cleanup-stale-upsert-operation", TaskID: "cleanup-stale-upsert-task", + CurrentSessionID: "cleanup-stale-upsert-old", + Attachments: []MessageAttachment{{AttachmentID: "cleanup-stale-upsert-attachment"}}, + } + if err := repo.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + const destinationSessionID = "cleanup-stale-upsert-new" + if err := repo.TransferSession(ctx, cleanup.SessionID, destinationSessionID); err != nil { + t.Fatal(err) + } + if err := repo.UpsertAttachmentCleanup(ctx, cleanup); err != nil { + t.Fatal(err) + } + stored, err := repo.GetAttachmentCleanup( + ctx, cleanup.SessionID, cleanup.EntryID, cleanup.OperationID, + ) + if err != nil { + t.Fatal(err) + } + if stored == nil || stored.CurrentSessionID != destinationSessionID { + t.Fatalf("cleanup after stale upsert = %#v", stored) + } +} + +func TestSQLiteAttachmentCleanupUpsertReadsTransferredQueueEntry(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + entry := insertTestEntry(t, repo, "cleanup-entry-old", "cleanup-entry-task", "prompt", QueuedByUser, nil, nil) + const destinationSessionID = "cleanup-entry-new" + if err := repo.TransferSession(ctx, entry.SessionID, destinationSessionID); err != nil { + t.Fatal(err) + } + if err := repo.UpsertAttachmentCleanup(ctx, AttachmentCleanup{ + SessionID: entry.SessionID, EntryID: entry.ID, OperationID: "cleanup-entry-operation", + TaskID: entry.TaskID, CurrentSessionID: entry.SessionID, + Attachments: []MessageAttachment{{AttachmentID: "cleanup-entry-attachment"}}, + }); err != nil { + t.Fatal(err) + } + stored, err := repo.GetAttachmentCleanup(ctx, entry.SessionID, entry.ID, "cleanup-entry-operation") + if err != nil { + t.Fatal(err) + } + if stored == nil || stored.CurrentSessionID != destinationSessionID { + t.Fatalf("cleanup after transferred queue entry = %#v", stored) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/repository_transfer_guard_regression_test.go b/apps/backend/internal/orchestrator/messagequeue/repository_transfer_guard_regression_test.go new file mode 100644 index 0000000000..68b6738957 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/repository_transfer_guard_regression_test.go @@ -0,0 +1,53 @@ +package messagequeue + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +// The exported transfer guard is the cross-package fence task repositories +// rely on: while a durable compensation row exists, attachment +// release/transfer/cleanup mutations on either session must fail closed. +func TestGuardSessionTransferRejectsActiveTransfer(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + sqlRepo := repo.(*sqliteRepository) + persistence := repo.(interface { + UpsertSessionTransferCompensation(context.Context, SessionTransferCompensation) error + DeleteSessionTransferCompensation(context.Context, string, string, string, string) error + }) + require.NoError(t, persistence.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "guard-regression-op", + TaskID: "guard-regression-task", FromSessionID: "guard-old", ToSessionID: "guard-new", + })) + + guard := func(sessionID string) error { + tx, err := sqlRepo.db.BeginTxx(ctx, nil) + require.NoError(t, err) + defer func() { _ = tx.Rollback() }() + // LockSessionInTransaction embeds the transfer guard, so the lock + // itself must fail closed on either transfer session. + return LockSessionInTransaction(ctx, tx, sqlRepo.db, sessionID) + } + require.ErrorIs(t, guard("guard-old"), ErrSessionTransferInProgress) + require.ErrorIs(t, guard("guard-new"), ErrSessionTransferInProgress) + require.NoError(t, guard("guard-unrelated")) + + // The exported guard also fails closed when the caller holds the raw + // session lock without going through LockSessionInTransaction. + func() { + tx, err := sqlRepo.db.BeginTxx(ctx, nil) + require.NoError(t, err) + defer func() { _ = tx.Rollback() }() + require.NoError(t, lockSessionTxIn(ctx, tx, sqlRepo.db, "guard-old")) + require.ErrorIs(t, GuardSessionTransferInTransaction(ctx, tx, sqlRepo.db, "guard-old"), ErrSessionTransferInProgress) + }() + + require.NoError(t, persistence.DeleteSessionTransferCompensation( + ctx, "guard-regression-op", "guard-regression-task", "guard-old", "guard-new", + )) + require.NoError(t, guard("guard-old")) + require.NoError(t, guard("guard-new")) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/requeue_fifo_regression_test.go b/apps/backend/internal/orchestrator/messagequeue/requeue_fifo_regression_test.go index c1b8a431ad..b2a99f9347 100644 --- a/apps/backend/internal/orchestrator/messagequeue/requeue_fifo_regression_test.go +++ b/apps/backend/internal/orchestrator/messagequeue/requeue_fifo_regression_test.go @@ -5,9 +5,23 @@ import ( "testing" "time" + "github.com/kandev/kandev/internal/common/logger" "github.com/stretchr/testify/require" ) +func newRequeueTestService(t *testing.T, repo Repository) *Service { + t.Helper() + log, err := logger.NewLogger(logger.LoggingConfig{ + Level: "error", + Format: "console", + OutputPath: "stderr", + }) + require.NoError(t, err) + service := NewService(repo, 10, log) + service.SetAutoMergeEnabled(false) + return service +} + type requeueRepositoryFactory struct { name string new func(*testing.T) Repository @@ -77,6 +91,7 @@ func TestRepository_RequeuePreservingFIFO_ReplacesPendingCoalesceTarget(t *testi insertRequeueTestEntry(t, repo, "session", "noise", QueuedByUser, nil) pending := insertRequeueTestEntry(t, repo, "session", "pending", QueuedByWorkflow, metadata) + pendingPosition := pending.Position require.NoError(t, repo.RequeuePreservingFIFO(ctx, taken)) @@ -90,6 +105,46 @@ func TestRepository_RequeuePreservingFIFO_ReplacesPendingCoalesceTarget(t *testi }) } } +func TestService_RequeueAtHeadInvalidatesLeasedCoalesceTarget(t *testing.T) { + factories := []requeueRepositoryFactory{ + {name: "memory", new: func(*testing.T) Repository { return NewMemoryRepository() }}, + {name: "sqlite", new: newTestSQLiteRepo}, + } + for _, factory := range factories { + t.Run(factory.name, func(t *testing.T) { + svc := newRequeueTestService(t, factory.new(t)) + ctx := context.Background() + metadata := map[string]interface{}{MetadataCoalesceKey: "retry-key"} + target, _, err := svc.QueueMessageWithCoalesceKey( + ctx, "session", "task", "original", "", QueuedByUser, false, + nil, metadata, "retry-key", true, + ) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, target.SessionID, target.ID, "connection") + require.NoError(t, err) + + retry := &QueuedMessage{ + SessionID: target.SessionID, + TaskID: target.TaskID, + Content: "retry", + QueuedBy: QueuedByUser, + Metadata: metadata, + } + require.NoError(t, svc.RequeueAtHead(ctx, retry)) + status := svc.GetStatus(ctx, target.SessionID) + require.Len(t, status.Entries, 1) + require.Equal(t, target.ID, status.Entries[0].ID) + require.Equal(t, "retry", status.Entries[0].Content) + + _, err = svc.UpdateMessageWithLease( + ctx, target.SessionID, target.ID, lease.LeaseID, + "stale-operation", "connection", lease.TargetRevision, + "stale edit", nil, nil, + ) + require.ErrorIs(t, err, ErrEditLeaseNotFound) + }) + } +} func TestRepository_RequeuePreservingFIFO_EmptyQueueStartsAtOne(t *testing.T) { for _, factory := range requeueRepositoryFactories() { @@ -162,3 +217,50 @@ func TestService_RequeueAtHeadWaitsForSessionAdmission(t *testing.T) { t.Fatal("requeue did not reach the repository after admission was released") } } +func TestRepository_RequeuePreservingFIFO_PersistsShiftedHighWater(t *testing.T) { + for _, factory := range autoRunRepositoryFactories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + ctx := context.Background() + entries := make([]*QueuedMessage, 0, 10) + for range 10 { + entries = append(entries, insertRequeueTestEntry(t, repo, "requeue-high-water", "entry", QueuedByUser, nil)) + } + taken, err := repo.TakeByID(ctx, "requeue-high-water", entries[1].ID) + require.NoError(t, err) + require.NotNil(t, taken) + require.NoError(t, repo.RequeuePreservingFIFO(ctx, taken)) + + for { + head, err := repo.TakeHead(ctx, "requeue-high-water") + require.NoError(t, err) + if head == nil { + break + } + } + later := insertRequeueTestEntry(t, repo, "requeue-high-water", "later", QueuedByUser, nil) + require.Greater(t, later.Position, int64(11)) + }) + } +} +func TestRepository_TransferDrainedSourceResetsPositionCounter(t *testing.T) { + for _, factory := range autoRunRepositoryFactories { + t.Run(factory.name, func(t *testing.T) { + repo := factory.new(t) + ctx := context.Background() + err := repo.Insert(ctx, &QueuedMessage{ + SessionID: "transfer-drained-source", + TaskID: "task", + Content: "entry", + QueuedBy: QueuedByUser, + }, 0) + require.NoError(t, err) + _, err = repo.TakeHead(ctx, "transfer-drained-source") + require.NoError(t, err) + require.NoError(t, repo.TransferSession(ctx, "transfer-drained-source", "transfer-destination")) + + reused := insertRequeueTestEntry(t, repo, "transfer-drained-source", "reused", QueuedByUser, nil) + require.Equal(t, int64(1), reused.Position) + }) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/reservation_generation_test.go b/apps/backend/internal/orchestrator/messagequeue/reservation_generation_test.go new file mode 100644 index 0000000000..412616747d --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/reservation_generation_test.go @@ -0,0 +1,37 @@ +package messagequeue + +import ( + "context" + "errors" + "testing" +) + +type generationReadFailureRepository struct { + Repository +} + +func (r *generationReadFailureRepository) SessionGeneration(context.Context, string) (int64, error) { + return 0, errors.New("non-transactional session generation read") +} + +func (r *generationReadFailureRepository) LifecycleGeneration(context.Context, string) (int64, error) { + return 0, errors.New("non-transactional lifecycle generation read") +} + +func TestReserveQueuedUsesRepositoryAtomicGenerationSnapshot(t *testing.T) { + ctx := context.Background() + repo := &generationReadFailureRepository{Repository: NewMemoryRepository()} + service := setupService(t) + service.repo = repo + source, err := service.QueueMessage(ctx, "session-atomic", "task-atomic", "prompt", "", QueuedByUser, false, nil) + if err != nil { + t.Fatal(err) + } + reserved, exists, autoRun := service.ReserveQueuedWithAutoRun(ctx, source.SessionID) + if !exists || !autoRun || reserved == nil || reserved.ID != source.ID { + t.Fatalf("atomic reservation = %#v, exists=%t, autoRun=%t", reserved, exists, autoRun) + } + if !reserved.reservationGenerationsCaptured { + t.Fatal("repository reservation did not capture fencing generations") + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/round78_regression_test.go b/apps/backend/internal/orchestrator/messagequeue/round78_regression_test.go new file mode 100644 index 0000000000..c87e2ae9ae --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/round78_regression_test.go @@ -0,0 +1,65 @@ +package messagequeue + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +type failAfterReservationRepository struct { + Repository + reserved bool +} + +func (r *failAfterReservationRepository) ReserveHeadIfAutoRun(ctx context.Context, sessionID string) (*QueuedMessage, bool, error) { + msg, enabled, err := r.Repository.ReserveHeadIfAutoRun(ctx, sessionID) + if err == nil && enabled && msg != nil { + r.reserved = true + } + return msg, enabled, err +} + +func (r *failAfterReservationRepository) LifecycleGeneration(ctx context.Context, taskID string) (int64, error) { + if r.reserved { + return 0, errors.New("lifecycle generation unavailable after reservation") + } + return r.Repository.LifecycleGeneration(ctx, taskID) +} + +func TestReserveQueuedWithAutoRunCapturesLifecycleGenerationBeforeReservation(t *testing.T) { + base := NewMemoryRepository() + repo := &failAfterReservationRepository{Repository: base} + svc := setupService(t) + svc.repo = repo + ctx := context.Background() + + _, _, accepted, err := svc.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-reserve-generation", "task-reserve-generation", "lifecycle", "", + QueuedByWorkflow, false, nil, nil, "lifecycle-key", true, + ) + require.NoError(t, err) + require.True(t, accepted) + + reserved, ok, autoRun := svc.ReserveQueuedWithAutoRun(ctx, "session-reserve-generation") + require.True(t, autoRun) + require.True(t, ok, "a successful reservation must be returned even when later reads fail") + require.NotNil(t, reserved) + require.Equal(t, "task-reserve-generation", reserved.TaskID) +} + +func TestEditLeaseDoesNotBlockAppendByAnotherQueueOwner(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session-owner-fencing", "task-owner-fencing", "user entry", "", QueuedByUser, false, nil) + require.NoError(t, err) + _, err = svc.BeginEdit(ctx, entry.SessionID, entry.ID, "connection-a") + require.NoError(t, err) + + inserted, appended, err := svc.AppendContent(ctx, entry.SessionID, entry.TaskID, "agent entry", "", QueuedByAgent, false, nil) + require.NoError(t, err) + require.False(t, appended) + require.NotEqual(t, entry.ID, inserted.ID) + require.Equal(t, QueuedByAgent, inserted.QueuedBy) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/round79_regression_test.go b/apps/backend/internal/orchestrator/messagequeue/round79_regression_test.go new file mode 100644 index 0000000000..27ebcec6f0 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/round79_regression_test.go @@ -0,0 +1,41 @@ +package messagequeue + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestReserveQueuedWithAutoRunCapturesReservedHeadGeneration(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + + _, err := svc.PurgeTask(ctx, "task-reserved-head") + require.NoError(t, err) + _, _, accepted, err := svc.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-reserved-head", "task-reserved-head", "head", "", + QueuedByWorkflow, false, nil, nil, "head-key", true, + ) + require.NoError(t, err) + require.True(t, accepted) + reserved, ok := svc.ReserveQueued(ctx, "session-reserved-head") + require.True(t, ok) + require.NotNil(t, reserved) + + _, _, accepted, err = svc.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-reserved-head", "task-following", "following", "", + QueuedByWorkflow, false, nil, nil, "following-key", true, + ) + require.NoError(t, err) + require.True(t, accepted) + + recovered, ok, autoRun := svc.ReserveQueuedWithAutoRun(ctx, "session-reserved-head") + require.True(t, autoRun) + require.True(t, ok) + require.NotNil(t, recovered) + require.Equal(t, "task-reserved-head", recovered.TaskID) + _, generation, captured := recovered.ReservationGenerations() + require.True(t, captured) + require.Equal(t, int64(1), generation) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/send_now.go b/apps/backend/internal/orchestrator/messagequeue/send_now.go index 6cede26bca..eb5aed373b 100644 --- a/apps/backend/internal/orchestrator/messagequeue/send_now.go +++ b/apps/backend/internal/orchestrator/messagequeue/send_now.go @@ -46,11 +46,23 @@ var ( // at its original position and every durable lifecycle row can be acknowledged // only after the replacement prompt is accepted. type SendNowClaim struct { + ClaimID string `json:"claim_id,omitempty"` Identity QueueSessionIdentity `json:"identity"` OperationGeneration int64 `json:"operation_generation"` Sources []QueuedMessage `json:"sources"` Dispatch QueuedMessage `json:"dispatch"` SourceGenerations map[string]int64 `json:"source_generations,omitempty"` + // SessionGeneration fences a restore after the session queue was purged. + // Ordinary sources are absent from storage while claimed, so task + // generations alone cannot prevent them from being resurrected. + SessionGeneration int64 `json:"session_generation"` +} + +// PendingSendNowClaim carries the durable acceptance decision used to choose +// between restoring and acknowledging an interrupted replacement dispatch. +type PendingSendNowClaim struct { + Claim SendNowClaim + Accepted bool } func sendNowSourceGenerationChanged(claim *SendNowClaim, source QueuedMessage, current int64) bool { @@ -61,6 +73,18 @@ func sendNowSourceGenerationChanged(claim *SendNowClaim, source QueuedMessage, c return ok && expected != current } +func sendNowClaimSourcesAllInvalidated(claim *SendNowClaim, generations map[string]int64) bool { + if claim == nil || len(claim.Sources) == 0 { + return false + } + for _, source := range claim.Sources { + if source.TaskID == "" || !sendNowSourceGenerationChanged(claim, source, generations[source.TaskID]) { + return false + } + } + return true +} + // ValidateSendNowEntries checks aggregate admission limits without building a // dispatch envelope. Callers use it before interrupting an active turn; the // repository still validates the exact snapshot again inside its claim diff --git a/apps/backend/internal/orchestrator/messagequeue/service.go b/apps/backend/internal/orchestrator/messagequeue/service.go index 431fc610c9..6557dbc542 100644 --- a/apps/backend/internal/orchestrator/messagequeue/service.go +++ b/apps/backend/internal/orchestrator/messagequeue/service.go @@ -11,8 +11,11 @@ package messagequeue import ( "context" + "crypto/sha256" + "encoding/json" "errors" "fmt" + "reflect" "sync" "sync/atomic" "time" @@ -35,8 +38,17 @@ type Service struct { autoMergeLoader atomic.Pointer[autoMergePolicyLoader] statusEpoch string logger *logger.Logger - admissionMu sync.Mutex - admissions map[string]*sessionAdmission + // lifecycleMu fences task-wide purges against every session-scoped queue + // admission. A task purge cannot enumerate every possibly empty session, + // so the barrier must cover admissions before they reach the repository. + lifecycleMu sync.RWMutex + admissionMu sync.Mutex + admissions map[string]*sessionAdmission + editLeaseMu sync.Mutex + editLeases map[editLeaseKey]*QueueEditLease + editRevisions map[editLeaseKey]int64 + editRevisionTaskIDs map[editLeaseKey]string + sessionTransferCompensationLeaseRenewInterval time.Duration } type sessionAdmission struct { @@ -47,8 +59,9 @@ type sessionAdmission struct { type sessionAdmissionContextKey struct{} type sessionAdmissionToken struct { - service *Service - sessionID string + service *Service + sessionID string + afterLifecycleUnlock *[]func() } type queueAttachmentRepository interface { @@ -90,10 +103,13 @@ type autoMergeAdmissionRepository interface { // pass 0 to disable the cap. func NewService(repo Repository, maxPerSession int, log *logger.Logger) *Service { service := &Service{ - repo: repo, - statusEpoch: uuid.NewString(), - logger: log.WithFields(zap.String("component", "message-queue")), - admissions: make(map[string]*sessionAdmission), + repo: repo, + statusEpoch: uuid.NewString(), + logger: log.WithFields(zap.String("component", "message-queue")), + admissions: make(map[string]*sessionAdmission), + editLeases: make(map[editLeaseKey]*QueueEditLease), + editRevisions: make(map[editLeaseKey]int64), + editRevisionTaskIDs: make(map[editLeaseKey]string), } service.SetMaxPerSession(maxPerSession) service.mergeEnabled.Store(true) @@ -136,6 +152,12 @@ func (s *Service) ListDurableLifecycleEntries(ctx context.Context) ([]QueuedMess return s.repo.ListDurableLifecycleEntries(ctx) } +// SessionGeneration returns the session destructive-mutation generation used +// to fence Send Now restores. +func (s *Service) SessionGeneration(ctx context.Context, sessionID string) (int64, error) { + return s.repo.SessionGeneration(ctx, sessionID) +} + // SetMaxPerSession applies a new admission cap without pruning existing rows. // Non-positive values disable the cap. func (s *Service) SetMaxPerSession(maxPerSession int) { @@ -266,10 +288,30 @@ func (s *Service) WithSessionAdmission(ctx context.Context, sessionID string, fn return errors.New("session admission callback is nil") } if token, ok := ctx.Value(sessionAdmissionContextKey{}).(sessionAdmissionToken); ok && - token.service == s && token.sessionID == sessionID { - return fn(ctx) + token.service == s { + if token.sessionID == sessionID { + return fn(ctx) + } + return s.withSessionAdmissionLock(ctx, sessionID, token.afterLifecycleUnlock, fn) } + var afterLifecycleUnlock []func() + s.lifecycleMu.RLock() + defer func() { + s.lifecycleMu.RUnlock() + for _, callback := range afterLifecycleUnlock { + callback() + } + }() + return s.withSessionAdmissionLock(ctx, sessionID, &afterLifecycleUnlock, fn) +} + +func (s *Service) withSessionAdmissionLock( + ctx context.Context, + sessionID string, + afterLifecycleUnlock *[]func(), + fn func(context.Context) error, +) error { s.admissionMu.Lock() entry := s.admissions[sessionID] if entry == nil { @@ -291,12 +333,664 @@ func (s *Service) WithSessionAdmission(ctx context.Context, sessionID string, fn }() admittedCtx := context.WithValue(ctx, sessionAdmissionContextKey{}, sessionAdmissionToken{ - service: s, - sessionID: sessionID, + service: s, sessionID: sessionID, afterLifecycleUnlock: afterLifecycleUnlock, }) return fn(admittedCtx) } +// withSessionAdmissions holds both queue admission locks in stable order. +// Transfer and replacement operations otherwise allow opposite-direction +// transfers to deadlock while each session is waiting for the other lock. +func (s *Service) withSessionAdmissions( + ctx context.Context, + firstSessionID, secondSessionID string, + fn func(context.Context) error, +) error { + if firstSessionID == secondSessionID { + return s.WithSessionAdmission(ctx, firstSessionID, fn) + } + if firstSessionID > secondSessionID { + firstSessionID, secondSessionID = secondSessionID, firstSessionID + } + return s.WithSessionAdmission(ctx, firstSessionID, func(admittedCtx context.Context) error { + return s.WithSessionAdmission(admittedCtx, secondSessionID, fn) + }) +} + +type editLeaseKey struct { + sessionID string + entryID string +} + +func (s *Service) editLeaseKey(sessionID, entryID string) editLeaseKey { + return editLeaseKey{sessionID: sessionID, entryID: entryID} +} + +func editOperationHash(content string, attachments []MessageAttachment, metadata map[string]interface{}) string { + payload, _ := json.Marshal(struct { + Content string + Attachments []MessageAttachment + Metadata map[string]interface{} + }{content, attachments, metadata}) + return fmt.Sprintf("%x", sha256.Sum256(payload)) +} + +func (s *Service) expireEditLeaseLocked(key editLeaseKey, now time.Time) { + if lease := s.editLeases[key]; lease != nil && !lease.ExpiresAt.After(now) { + delete(s.editLeases, key) + } +} + +type sessionTransferFenceRepository interface { + withSessionTransferFence(context.Context, string, func(context.Context) error) error +} + +type editReplayValidationRepository interface { + validateSessionEntryForEditReplay(context.Context, string, string) error +} + +type editLeaseRepository interface { + acquireEditLease(context.Context, *QueueEditLease) error + renewEditLease(context.Context, *QueueEditLease) error + releaseEditLease(context.Context, string, string, string) error +} + +func (s *Service) validateDuplicateEditReplay( + ctx context.Context, + sessionID, entryID string, +) error { + repo, ok := s.repo.(editReplayValidationRepository) + if !ok { + _, err := s.findQueuedMessageForEdit(ctx, sessionID, entryID) + return err + } + return repo.validateSessionEntryForEditReplay(ctx, sessionID, entryID) +} + +func (s *Service) withRepositorySessionTransferFence( + ctx context.Context, + sessionID string, + fn func(context.Context) error, +) error { + repo, ok := s.repo.(sessionTransferFenceRepository) + if !ok { + return fn(ctx) + } + return repo.withSessionTransferFence(ctx, sessionID, fn) +} + +func (s *Service) editableEntry( + ctx context.Context, + sessionID, entryID string, +) (*QueuedMessage, error) { + entries, err := s.repo.ListBySession(ctx, sessionID) + if err != nil { + return nil, err + } + for i := range entries { + if entries[i].ID == entryID && + entries[i].QueuedBy == QueuedByUser && + !entries[i].IsReservedInFlight() { + return &entries[i], nil + } + } + return nil, ErrEditLeaseNotFound +} + +// BeginEdit acquires a target-bound lease for a visible user-owned entry. +func (s *Service) BeginEdit(ctx context.Context, sessionID, entryID, connectionID string) (*QueueEditLease, error) { + if sessionID == "" || entryID == "" || connectionID == "" { + return nil, ErrEditLeaseNotFound + } + var lease *QueueEditLease + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + entry, err := s.editableEntry(admittedCtx, sessionID, entryID) + if err != nil { + return err + } + key := s.editLeaseKey(sessionID, entryID) + now := time.Now().UTC() + s.editLeaseMu.Lock() + s.expireEditLeaseLocked(key, now) + if s.editLeases[key] != nil { + s.editLeaseMu.Unlock() + return ErrEditConflict + } + lease = &QueueEditLease{ + SessionID: sessionID, EntryID: entryID, LeaseID: uuid.NewString(), + TargetRevision: s.editRevisions[key], LeaseGeneration: 1, + ExpiresAt: now.Add(QueueEditLeaseTTL), connectionID: connectionID, + taskID: entry.TaskID, + } + s.editLeaseMu.Unlock() + if repo, ok := s.repo.(editLeaseRepository); ok { + if err := repo.acquireEditLease(admittedCtx, lease); err != nil { + return err + } + } + s.editLeaseMu.Lock() + s.editLeases[key] = lease + s.editLeaseMu.Unlock() + return nil + }) + if err != nil { + return nil, err + } + return cloneEditLease(lease), nil +} + +// GetEditLease returns the current lease for an entry without acquiring +// session admission. Callers use it only to avoid starting cleanup while an +// editor still owns the target; the mutation itself remains admission-bound. +func (s *Service) GetEditLease(_ context.Context, sessionID, entryID string) (*QueueEditLease, error) { + key := s.editLeaseKey(sessionID, entryID) + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.expireEditLeaseLocked(key, time.Now().UTC()) + lease := s.editLeases[key] + if lease == nil { + return nil, ErrEditLeaseNotFound + } + return cloneEditLease(lease), nil +} + +// RenewEdit extends a live lease owned by connectionID. +func (s *Service) RenewEdit(ctx context.Context, sessionID, entryID, leaseID, connectionID string) (*QueueEditLease, error) { + var renewed *QueueEditLease + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + if _, err := s.editableEntry(admittedCtx, sessionID, entryID); err != nil { + return err + } + return s.withRepositorySessionTransferFence(admittedCtx, sessionID, func(context.Context) error { + key := s.editLeaseKey(sessionID, entryID) + now := time.Now().UTC() + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.expireEditLeaseLocked(key, now) + lease := s.editLeases[key] + if lease == nil || lease.LeaseID != leaseID || leaseConnection(lease) != connectionID { + return ErrEditLeaseNotFound + } + lease.LeaseGeneration++ + lease.ExpiresAt = now.Add(QueueEditLeaseTTL) + renewed = cloneEditLease(lease) + return nil + }) + }) + if err != nil || renewed == nil { + return renewed, err + } + if repo, ok := s.repo.(editLeaseRepository); ok { + if err := repo.renewEditLease(ctx, renewed); err != nil { + s.editLeaseMu.Lock() + if lease := s.editLeases[s.editLeaseKey(sessionID, entryID)]; lease != nil && lease.LeaseID == leaseID { + delete(s.editLeases, s.editLeaseKey(sessionID, entryID)) + } + s.editLeaseMu.Unlock() + return nil, err + } + } + return renewed, nil +} + +func (s *Service) EndEdit(ctx context.Context, sessionID, entryID, leaseID, connectionID string) error { + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + return s.withRepositorySessionTransferFence(admittedCtx, sessionID, func(context.Context) error { + key := s.editLeaseKey(sessionID, entryID) + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.expireEditLeaseLocked(key, time.Now().UTC()) + lease := s.editLeases[key] + if lease == nil || lease.LeaseID != leaseID || leaseConnection(lease) != connectionID { + return ErrEditLeaseNotFound + } + delete(s.editLeases, key) + return nil + }) + }) + if err != nil { + return err + } + if repo, ok := s.repo.(editLeaseRepository); ok { + return repo.releaseEditLease(ctx, sessionID, entryID, leaseID) + } + return nil +} + +// EndEditAfterSave releases a lease and reports whether its latest update +// completed all attachment finalization. Only that state may authorize a +// post-save automatic drain. +func (s *Service) EndEditAfterSave(ctx context.Context, sessionID, entryID, leaseID, connectionID string) (bool, error) { + saved := false + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + return s.withRepositorySessionTransferFence(admittedCtx, sessionID, func(context.Context) error { + key := s.editLeaseKey(sessionID, entryID) + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.expireEditLeaseLocked(key, time.Now().UTC()) + lease := s.editLeases[key] + if lease == nil || lease.LeaseID != leaseID || leaseConnection(lease) != connectionID { + return ErrEditLeaseNotFound + } + saved = lease.lastOperationID != "" && lease.lastOperationFinalized + delete(s.editLeases, key) + return nil + }) + }) + if err != nil { + return false, err + } + if repo, ok := s.repo.(editLeaseRepository); ok { + if err := repo.releaseEditLease(ctx, sessionID, entryID, leaseID); err != nil { + return saved, err + } + } + return saved, nil +} + +func (s *Service) UpdateMessageWithLease( + ctx context.Context, + sessionID, entryID, leaseID, operationID, connectionID string, + expectedRevision int64, + content string, + attachments []MessageAttachment, + metadataUpdates map[string]interface{}, +) (int64, error) { + return s.UpdateMessageWithLeaseAfterValidation( + ctx, sessionID, entryID, leaseID, operationID, connectionID, + expectedRevision, content, attachments, metadataUpdates, nil, nil, + ) +} + +// UpdateMessageWithLeaseAfterValidation runs prepare after all rejectable edit +// preconditions pass and before the queue row is changed. Both callbacks run +// inside the session admission boundary, so rejected prepared state can be +// rolled back before another editor can acquire the target. +func (s *Service) UpdateMessageWithLeaseAfterValidation( + ctx context.Context, + sessionID, entryID, leaseID, operationID, connectionID string, + expectedRevision int64, + content string, + attachments []MessageAttachment, + metadataUpdates map[string]interface{}, + prepare func(context.Context) error, + rollback func(context.Context) error, +) (int64, error) { + return s.updateMessageWithLeaseAfterValidation( + ctx, sessionID, entryID, leaseID, operationID, connectionID, + expectedRevision, content, attachments, metadataUpdates, prepare, rollback, nil, + ) +} + +// UpdateMessageWithLeaseAfterValidationAndFinalize is the attachment-aware +// variant. finalize runs after the durable row update while session admission +// remains held, preventing a concurrent session transfer from moving the row +// before superseded external state is reconciled. A failed finalize is retained +// as part of the operation record so an identical retry can run it again with +// the original pre-update snapshot. Intervening entry lifecycle mutations +// invalidate the lease operation record, so a valid replay cannot belong to a +// newer edit. +func (s *Service) UpdateMessageWithLeaseAfterValidationAndFinalize( + ctx context.Context, + sessionID, entryID, leaseID, operationID, connectionID string, + expectedRevision int64, + content string, + attachments []MessageAttachment, + metadataUpdates map[string]interface{}, + prepare func(context.Context) error, + rollback func(context.Context) error, + finalize func(context.Context, *QueuedMessage) error, +) (int64, error) { + return s.updateMessageWithLeaseAfterValidation( + ctx, sessionID, entryID, leaseID, operationID, connectionID, + expectedRevision, content, attachments, metadataUpdates, prepare, rollback, finalize, + ) +} + +func (s *Service) updateMessageWithLeaseAfterValidation( + ctx context.Context, + sessionID, entryID, leaseID, operationID, connectionID string, + expectedRevision int64, + content string, + attachments []MessageAttachment, + metadataUpdates map[string]interface{}, + prepare func(context.Context) error, + rollback func(context.Context) error, + finalize func(context.Context, *QueuedMessage) error, +) (int64, error) { + if operationID == "" { + return 0, ErrEditLeaseNotFound + } + key := s.editLeaseKey(sessionID, entryID) + operationHash := editOperationHash(content, attachments, metadataUpdates) + var revision int64 + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + initialRevision, initialOperationID, duplicate, previous, err := s.beginEditMutation( + key, leaseID, operationID, connectionID, expectedRevision, operationHash, + ) + if err != nil { + return err + } + if duplicate { + if err := s.validateDuplicateEditReplay(admittedCtx, sessionID, entryID); err != nil { + return err + } + revision = initialRevision + return s.finalizeDuplicateEdit( + admittedCtx, key, operationID, operationHash, previous, finalize, + ) + } + revision = initialRevision + previous, err = s.findQueuedMessageForEdit(admittedCtx, sessionID, entryID) + if err != nil { + return err + } + + if prepare != nil { + if err := prepare(admittedCtx); err != nil { + s.rollbackPreparedEdit(admittedCtx, rollback) + return err + } + } + + lease, err := s.relockEditMutation(key, leaseID, connectionID, revision, initialOperationID) + if err != nil { + s.rollbackPreparedEdit(admittedCtx, rollback) + return err + } + if leaseRepository, ok := s.repo.(durableEditLeaseMutationRepository); ok { + if err := leaseRepository.updateContentAndMetadataWithLease(admittedCtx, sessionID, entryID, leaseID, content, attachments, metadataUpdates, QueuedByUser); err != nil { + s.editLeaseMu.Unlock() + s.rollbackPreparedEdit(admittedCtx, rollback) + return err + } + } else if err := s.repo.UpdateContentAndMetadata(admittedCtx, sessionID, entryID, content, attachments, metadataUpdates, QueuedByUser); err != nil { + s.editLeaseMu.Unlock() + s.rollbackPreparedEdit(admittedCtx, rollback) + return err + } + revision++ + s.editRevisions[key] = revision + s.editRevisionTaskIDs[key] = lease.taskID + lease.TargetRevision = revision + lease.lastOperationID = operationID + lease.lastOperationHash = operationHash + lease.lastOperationResult = revision + lease.lastOperationPrevious = cloneQueuedMessage(previous) + lease.lastOperationFinalized = finalize == nil + s.editLeaseMu.Unlock() + if finalize != nil { + if err := finalize(admittedCtx, previous); err != nil { + return err + } + s.markEditOperationFinalized(key, operationID, operationHash) + } + return nil + }) + return revision, err +} + +func (s *Service) finalizeDuplicateEdit( + ctx context.Context, + key editLeaseKey, + operationID, operationHash string, + previous *QueuedMessage, + finalize func(context.Context, *QueuedMessage) error, +) error { + if finalize == nil || s.editOperationFinalized(key, operationID, operationHash) { + return nil + } + if err := finalize(ctx, previous); err != nil { + return err + } + s.markEditOperationFinalized(key, operationID, operationHash) + return nil +} + +func (s *Service) findQueuedMessageForEdit( + ctx context.Context, + sessionID, entryID string, +) (*QueuedMessage, error) { + entries, err := s.repo.ListBySession(ctx, sessionID) + if err != nil { + return nil, err + } + for i := range entries { + if entries[i].ID == entryID { + return cloneQueuedMessage(&entries[i]), nil + } + } + return nil, ErrEntryNotFound +} +func (s *Service) beginEditMutation( + key editLeaseKey, + leaseID, operationID, connectionID string, + expectedRevision int64, + operationHash string, +) (int64, string, bool, *QueuedMessage, error) { + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.expireEditLeaseLocked(key, time.Now().UTC()) + lease := s.editLeases[key] + if lease == nil || lease.LeaseID != leaseID || leaseConnection(lease) != connectionID { + return 0, "", false, nil, ErrEditLeaseNotFound + } + if lease.lastOperationID == operationID { + if lease.lastOperationHash == operationHash { + return lease.lastOperationResult, lease.lastOperationID, true, + cloneQueuedMessage(lease.lastOperationPrevious), nil + } + return 0, "", false, nil, ErrEditRevisionConflict + } + revision := s.editRevisions[key] + if expectedRevision != revision { + return 0, "", false, nil, ErrEditRevisionConflict + } + return revision, lease.lastOperationID, false, nil, nil +} +func (s *Service) editOperationFinalized(key editLeaseKey, operationID, operationHash string) bool { + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + lease := s.editLeases[key] + return lease != nil && + lease.lastOperationID == operationID && + lease.lastOperationHash == operationHash && + lease.lastOperationFinalized +} + +func (s *Service) markEditOperationFinalized(key editLeaseKey, operationID, operationHash string) { + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + lease := s.editLeases[key] + if lease != nil && lease.lastOperationID == operationID && lease.lastOperationHash == operationHash { + lease.lastOperationFinalized = true + } +} + +// relockEditMutation validates the lease again after attachment preparation and +// leaves editLeaseMu held for the repository update. +func (s *Service) relockEditMutation( + key editLeaseKey, + leaseID, connectionID string, + revision int64, + initialOperationID string, +) (*QueueEditLease, error) { + s.editLeaseMu.Lock() + s.expireEditLeaseLocked(key, time.Now().UTC()) + lease := s.editLeases[key] + if lease == nil || lease.LeaseID != leaseID || leaseConnection(lease) != connectionID { + s.editLeaseMu.Unlock() + return nil, ErrEditLeaseNotFound + } + if s.editRevisions[key] != revision || lease.lastOperationID != initialOperationID { + s.editLeaseMu.Unlock() + return nil, ErrEditRevisionConflict + } + return lease, nil +} + +func (s *Service) rollbackPreparedEdit(ctx context.Context, rollback func(context.Context) error) { + if rollback == nil { + return + } + if err := rollback(ctx); err != nil { + s.logger.Warn("failed to roll back prepared queue edit state", zap.Error(err)) + } +} + +func cloneEditLease(lease *QueueEditLease) *QueueEditLease { + if lease == nil { + return nil + } + copy := *lease + return © +} + +// leaseConnection is intentionally process-local. The connection binding is +// stored beside the lease rather than exposed in the wire response. +func leaseConnection(lease *QueueEditLease) string { + return lease.connectionID +} + +func (s *Service) editLeaseBlocksHead(ctx context.Context, sessionID string) (bool, error) { + entries, err := s.repo.ListBySession(ctx, sessionID) + if err != nil { + return false, err + } + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + now := time.Now().UTC() + for _, entry := range entries { + key := s.editLeaseKey(sessionID, entry.ID) + s.expireEditLeaseLocked(key, now) + return s.editLeases[key] != nil, nil + } + return false, nil +} + +func (s *Service) editLeaseBlocksEntryLocked(sessionID, entryID string) bool { + key := s.editLeaseKey(sessionID, entryID) + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.expireEditLeaseLocked(key, time.Now().UTC()) + return s.editLeases[key] != nil +} +func (s *Service) editLeaseBlocksReorderLocked(sessionID string) bool { + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + now := time.Now().UTC() + for key := range s.editLeases { + if key.sessionID != sessionID { + continue + } + s.expireEditLeaseLocked(key, now) + if s.editLeases[key] != nil { + return true + } + } + return false +} + +func (s *Service) editLeaseBlocksTail(ctx context.Context, sessionID, excludedEntryID, queuedBy string) (bool, error) { + entries, err := s.repo.ListBySession(ctx, sessionID) + if err != nil { + return false, err + } + var tail *QueuedMessage + for i := range entries { + entry := &entries[i] + if entry.ID == excludedEntryID { + continue + } + if tail == nil || entry.Position > tail.Position { + tail = entry + } + } + if tail == nil { + return false, nil + } + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + key := s.editLeaseKey(sessionID, tail.ID) + s.expireEditLeaseLocked(key, time.Now().UTC()) + return s.editLeases[key] != nil && tail.QueuedBy == queuedBy, nil +} + +func (s *Service) editLeaseBlocksMerge(ctx context.Context, sessionID, sourceID string) (bool, error) { + entries, err := s.repo.ListBySession(ctx, sessionID) + if err != nil { + return false, err + } + var source, target *QueuedMessage + for i := range entries { + if entries[i].ID == sourceID { + source = &entries[i] + break + } + } + if source == nil { + return false, nil + } + for i := range entries { + entry := &entries[i] + if entry.IsReservedInFlight() || entry.Position >= source.Position { + continue + } + if target == nil || entry.Position > target.Position { + target = entry + } + } + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + now := time.Now().UTC() + sourceKey := s.editLeaseKey(sessionID, source.ID) + s.expireEditLeaseLocked(sourceKey, now) + if s.editLeases[sourceKey] != nil { + return true, nil + } + if target == nil { + return false, nil + } + targetKey := s.editLeaseKey(sessionID, target.ID) + s.expireEditLeaseLocked(targetKey, now) + return s.editLeases[targetKey] != nil, nil +} + +func (s *Service) deleteEditStateLocked(key editLeaseKey) { + delete(s.editLeases, key) + delete(s.editRevisions, key) + delete(s.editRevisionTaskIDs, key) +} + +func (s *Service) invalidateEditLeasesLocked(sessionIDs ...string) { + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + for key := range s.editRevisions { + for _, sessionID := range sessionIDs { + if key.sessionID == sessionID { + s.deleteEditStateLocked(key) + break + } + } + } + for key := range s.editLeases { + for _, sessionID := range sessionIDs { + if key.sessionID == sessionID { + s.deleteEditStateLocked(key) + break + } + } + } + for key := range s.editRevisionTaskIDs { + for _, sessionID := range sessionIDs { + if key.sessionID == sessionID { + s.deleteEditStateLocked(key) + break + } + } + } +} + // QueueMessage appends a new entry to the session's FIFO queue. Returns // ErrQueueFull when the cap is exceeded. func (s *Service) QueueMessage(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment) (*QueuedMessage, error) { @@ -309,6 +1003,14 @@ func (s *Service) QueueMessage(ctx context.Context, sessionID, taskID, content, func (s *Service) QueueMessageWithMetadata(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}) (*QueuedMessage, error) { identity, err := s.repo.ResolveSessionIdentity(ctx, taskID, sessionID) if err != nil { + // The legacy admission API predates durable session incarnations. Keep + // it usable for callers that only have the session/task pair while the + // identity-bound API remains strict for browser mutations. + if errors.Is(err, ErrSessionIdentityMismatch) { + return s.queueMessageWithMetadataAdmission( + ctx, nil, sessionID, taskID, content, model, userID, planMode, attachments, metadata, nil, nil, + ) + } return nil, err } return s.queueMessageWithMetadataAdmission( @@ -323,6 +1025,11 @@ func (s *Service) QueueMessageWithMetadata(ctx context.Context, sessionID, taskI func (s *Service) QueueMessageWithMetadataAfterInsert(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, afterInsert func(context.Context, *QueuedMessage) error) (*QueuedMessage, error) { identity, err := s.repo.ResolveSessionIdentity(ctx, taskID, sessionID) if err != nil { + if errors.Is(err, ErrSessionIdentityMismatch) { + return s.queueMessageWithMetadataAdmission( + ctx, nil, sessionID, taskID, content, model, userID, planMode, attachments, metadata, nil, afterInsert, + ) + } return nil, err } return s.queueMessageWithMetadataAdmission( @@ -544,6 +1251,17 @@ func (s *Service) finalizeAutoMerge( policy *AutoMergePolicy, ) *QueuedMessage { for policy != nil && policy.Enabled { + blocked, leaseErr := s.editLeaseBlocksTail(ctx, source.SessionID, source.ID, source.QueuedBy) + if leaseErr != nil { + s.logger.Error("automatic queue merge lease check failed; preserving separate admission", + zap.String("session_id", source.SessionID), + zap.String("source_entry_id", source.ID), + zap.Error(leaseErr)) + return source + } + if blocked { + return source + } var merged *QueuedMessage var didMerge bool var err error @@ -679,6 +1397,7 @@ func (s *Service) restoreMessage(ctx context.Context, msg *QueuedMessage) (*Queu if err := s.repo.Restore(ctx, &restored, 0); err != nil { return nil, err } + s.invalidateEditLease(restored.SessionID, restored.ID) s.logger.Info("message restored at original queue position", zap.String("session_id", restored.SessionID), zap.String("task_id", restored.TaskID), @@ -765,6 +1484,9 @@ func (s *Service) insertQueueMessageWithCoalesceKey(ctx context.Context, session } return nil, false, err } + if replaced && queued != nil { + s.invalidateEditLease(sessionID, queued.ID) + } s.logger.Info("message queued with coalesce key", zap.String("session_id", sessionID), zap.String("task_id", taskID), @@ -810,7 +1532,11 @@ func (s *Service) RequeueAtHead(ctx context.Context, msg *QueuedMessage) error { return errors.New("queued message is nil") } return s.WithSessionAdmission(ctx, msg.SessionID, func(admittedCtx context.Context) error { - return s.repo.RequeuePreservingFIFO(admittedCtx, msg) + if err := s.repo.RequeuePreservingFIFO(admittedCtx, msg); err != nil { + return err + } + s.invalidateEditLease(msg.SessionID, msg.ID) + return nil }) } @@ -1006,10 +1732,107 @@ func (s *Service) insertLifecycleMessageWithCoalesceKey(ctx context.Context, ide return queued, replaced, true, nil } +// InvalidateEditLeasesForTask drops process-local edit leases after the task +// repository has purged its durable queue rows in the same lifecycle operation. +func (s *Service) InvalidateEditLeasesForTask(taskID string) { + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.invalidateEditLeasesForTaskLocked(taskID) +} + +func (s *Service) invalidateEditLeasesForTaskLocked(taskID string) { + for key, lease := range s.editLeases { + if lease.taskID == taskID { + s.deleteEditStateLocked(key) + } + } + for key, revisionTaskID := range s.editRevisionTaskIDs { + if revisionTaskID == taskID { + s.deleteEditStateLocked(key) + } + } +} + +// InvalidateEditLeasesForSession drops process-local edit leases after a task +// repository deletes that session's durable queue rows. +func (s *Service) InvalidateEditLeasesForSession(sessionID string) { + s.invalidateEditLeasesLocked(sessionID) +} + +func (s *Service) invalidateEditLease(sessionID, entryID string) { + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + s.deleteEditStateLocked(s.editLeaseKey(sessionID, entryID)) +} + +// ReleaseEditLeasesForConnection drops every lease owned by a disconnected +// WebSocket connection. Durable release is attempted before local removal, +// but a failed durable release cannot preserve a disconnected owner's lease. +func (s *Service) ReleaseEditLeasesForConnection(connectionID string) int { + if connectionID == "" { + return 0 + } + s.editLeaseMu.Lock() + leases := make([]*QueueEditLease, 0) + for _, lease := range s.editLeases { + if lease.connectionID == connectionID { + leases = append(leases, cloneEditLease(lease)) + } + } + s.editLeaseMu.Unlock() + released := 0 + repo, persistent := s.repo.(editLeaseRepository) + for _, lease := range leases { + if persistent { + if err := repo.releaseEditLease(context.Background(), lease.SessionID, lease.EntryID, lease.LeaseID); err != nil { + s.logger.Warn("failed to release durable queue edit lease", + zap.String("session_id", lease.SessionID), + zap.String("entry_id", lease.EntryID), + zap.Error(err)) + } + } + s.editLeaseMu.Lock() + key := s.editLeaseKey(lease.SessionID, lease.EntryID) + if current := s.editLeases[key]; current != nil && current.LeaseID == lease.LeaseID { + delete(s.editLeases, key) + released++ + } + s.editLeaseMu.Unlock() + } + return released +} + // PurgeTask is a backend-only task lifecycle operation. Client deletion APIs // retain their reserved-entry protections. func (s *Service) PurgeTask(ctx context.Context, taskID string) (int, error) { - return s.repo.PurgeTask(ctx, taskID) + if token, ok := ctx.Value(sessionAdmissionContextKey{}).(sessionAdmissionToken); ok && + token.service == s && token.afterLifecycleUnlock != nil { + purgeCtx := context.WithoutCancel(ctx) + *token.afterLifecycleUnlock = append(*token.afterLifecycleUnlock, func() { + if _, err := s.purgeTask(purgeCtx, taskID); err != nil { + s.logger.Error("failed to purge task queue after admission", + zap.String("task_id", taskID), zap.Error(err)) + } + }) + return 0, nil + } + return s.purgeTask(ctx, taskID) +} + +func (s *Service) purgeTask(ctx context.Context, taskID string) (int, error) { + s.lifecycleMu.Lock() + defer s.lifecycleMu.Unlock() + s.editLeaseMu.Lock() + defer s.editLeaseMu.Unlock() + + // The lifecycle barrier above serializes the persistent purge with all + // queue admissions, including writes for sessions absent from the purge. + removed, err := s.repo.PurgeTask(ctx, taskID) + if err != nil { + return 0, err + } + s.invalidateEditLeasesForTaskLocked(taskID) + return removed, nil } // lifecycleGenerationFromMetadata reads the lifecycle generation captured on an entry. @@ -1046,7 +1869,14 @@ func (s *Service) ReserveQueuedWithAutoRun(ctx context.Context, sessionID string var msg *QueuedMessage autoRun := true err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { - var err error + blocked, err := s.editLeaseBlocksHead(admittedCtx, sessionID) + if err != nil { + return err + } + if blocked { + autoRun, err = s.repo.GetAutoRun(admittedCtx, sessionID) + return err + } msg, autoRun, err = s.repo.ReserveHeadIfAutoRun(admittedCtx, sessionID) return err }) @@ -1130,15 +1960,18 @@ func (s *Service) PauseAutoRunIfPendingForSession( return paused, err } -// AcknowledgeQueued removes a server-reserved entry after prompt acceptance. -func (s *Service) AcknowledgeQueued(ctx context.Context, sessionID, entryID string) error { - err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { - return s.repo.AcknowledgeByID(admittedCtx, sessionID, entryID) - }) - if errors.Is(err, ErrEntryNotFound) { - return nil +// AcknowledgeQueued settles only the exact server reservation carried by msg. +func (s *Service) AcknowledgeQueued(ctx context.Context, msg *QueuedMessage) error { + if msg == nil { + return errors.New("queued message is nil") } - return err + return s.WithSessionAdmission(ctx, msg.SessionID, func(admittedCtx context.Context) error { + err := s.repo.AcknowledgeReserved(admittedCtx, msg) + if err != nil && !errors.Is(err, ErrEntryNotFound) { + return err + } + return s.deletePendingQueueDispatch(admittedCtx, msg) + }) } func (s *Service) AcknowledgeQueuedForSession( @@ -1242,17 +2075,100 @@ func copyMessageMetadata(metadata map[string]interface{}, extraCapacity int) map } out := make(map[string]interface{}, len(metadata)+extraCapacity) for key, value := range metadata { - out[key] = value + out[key] = copyMessageMetadataValue(value) } return out } +func copyMessageMetadataValue(value interface{}) interface{} { + cloned := copyMessageMetadataReflect(reflect.ValueOf(value)) + if !cloned.IsValid() { + return nil + } + return cloned.Interface() +} + +func copyMessageMetadataReflect(value reflect.Value) reflect.Value { + if !value.IsValid() { + return value + } + switch value.Kind() { + case reflect.Interface: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + cloned := copyMessageMetadataReflect(value.Elem()) + out := reflect.New(value.Type()).Elem() + out.Set(cloned) + return out + case reflect.Map: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + out := reflect.MakeMapWithSize(value.Type(), value.Len()) + iter := value.MapRange() + for iter.Next() { + out.SetMapIndex(iter.Key(), copyMessageMetadataReflect(iter.Value())) + } + return out + case reflect.Slice: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + out := reflect.MakeSlice(value.Type(), value.Len(), value.Len()) + for i := range value.Len() { + out.Index(i).Set(copyMessageMetadataReflect(value.Index(i))) + } + return out + case reflect.Array: + out := reflect.New(value.Type()).Elem() + for i := range value.Len() { + out.Index(i).Set(copyMessageMetadataReflect(value.Index(i))) + } + return out + case reflect.Pointer: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + out := reflect.New(value.Type().Elem()) + out.Elem().Set(copyMessageMetadataReflect(value.Elem())) + return out + default: + return value + } +} + // AppendContent appends content onto the session's tail entry when the tail's // queued_by matches userID. Otherwise inserts a new entry. Returns // ErrQueueFull when an insert would exceed the cap. func (s *Service) AppendContent(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment) (*QueuedMessage, bool, error) { identity, err := s.repo.ResolveSessionIdentity(ctx, taskID, sessionID) if err != nil { + if errors.Is(err, ErrSessionIdentityMismatch) { + var msg *QueuedMessage + var appended bool + err = s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + blocked, blockErr := s.editLeaseBlocksTail(admittedCtx, sessionID, "", userID) + if blockErr != nil { + return blockErr + } + if blocked { + return ErrEditConflict + } + msg, appended, err = s.repo.AppendOrInsertTail( + admittedCtx, sessionID, taskID, content, model, userID, planMode, attachments, nil, s.MaxPerSession(), + ) + return err + }) + if err != nil { + return nil, false, err + } + s.logger.Info("queue append-or-insert", + zap.String("session_id", sessionID), + zap.String("entry_id", msg.ID), + zap.Bool("appended", appended)) + return msg, appended, nil + } return nil, false, err } return s.AppendContentForSession(ctx, identity, content, model, userID, planMode, attachments) @@ -1263,7 +2179,13 @@ func (s *Service) AppendContentForSession(ctx context.Context, identity QueueSes var msg *QueuedMessage var appended bool err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error { - var err error + blocked, err := s.editLeaseBlocksTail(admittedCtx, identity.SessionID, "", userID) + if err != nil { + return err + } + if blocked { + return ErrEditConflict + } msg, appended, err = s.repo.AppendOrInsertTailForSession(admittedCtx, identity, content, model, userID, planMode, attachments, nil, s.MaxPerSession()) return err }) @@ -1278,12 +2200,29 @@ func (s *Service) AppendContentForSession(ctx context.Context, identity QueueSes } // TakeQueued atomically removes and returns the head entry. Returns nil, false -// when the queue is empty. func (s *Service) TakeQueued(ctx context.Context, sessionID string) (*QueuedMessage, bool) { var msg *QueuedMessage err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { - var err error + // TakeQueued is a destructive legacy cleanup path. Never let it + // consume a durable lifecycle row that is already reserved for an + // in-flight dispatch; that row must remain available for ack/retry. + entries, err := s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + return err + } + if len(entries) > 0 && entries[0].IsReservedInFlight() { + return nil + } + blocked, err := s.editLeaseBlocksHead(admittedCtx, sessionID) + if err != nil || blocked { + return err + } msg, err = s.repo.TakeHead(admittedCtx, sessionID) + if err == nil && msg != nil { + s.editLeaseMu.Lock() + s.deleteEditStateLocked(s.editLeaseKey(sessionID, msg.ID)) + s.editLeaseMu.Unlock() + } return err }) if err != nil { @@ -1311,7 +2250,23 @@ func (s *Service) TakeQueuedIfAutoRun(ctx context.Context, sessionID string) (*Q if err != nil || !autoRun { return err } + entries, err := s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + return err + } + if len(entries) == 0 || entries[0].IsDurableLifecycle() || entries[0].IsReservedInFlight() { + return nil + } + blocked, err := s.editLeaseBlocksHead(admittedCtx, sessionID) + if err != nil || blocked { + return err + } msg, err = s.repo.TakeHead(admittedCtx, sessionID) + if err == nil && msg != nil { + s.editLeaseMu.Lock() + s.deleteEditStateLocked(s.editLeaseKey(sessionID, msg.ID)) + s.editLeaseMu.Unlock() + } return err }) if err != nil { @@ -1344,8 +2299,16 @@ func (s *Service) TakeQueuedIfAutoRun(ctx context.Context, sessionID string) (*Q func (s *Service) TakeQueuedEntry(ctx context.Context, sessionID, entryID string) (*QueuedMessage, bool, error) { var msg *QueuedMessage err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + if s.editLeaseBlocksEntryLocked(sessionID, entryID) { + return ErrEditConflict + } var err error msg, err = s.repo.TakeByID(admittedCtx, sessionID, entryID) + if err == nil && msg != nil { + s.editLeaseMu.Lock() + s.deleteEditStateLocked(s.editLeaseKey(sessionID, msg.ID)) + s.editLeaseMu.Unlock() + } return err }) if err != nil { @@ -1404,17 +2367,69 @@ func (s *Service) UpdateMessage(ctx context.Context, sessionID, entryID, content // lifecycle-aware callers that must inspect the trusted task ID and previous // attachment descriptors before replacing an entry. func (s *Service) GetEntry(ctx context.Context, sessionID, entryID string) (*QueuedMessage, error) { - entries, err := s.repo.ListBySession(ctx, sessionID) - if err != nil { + var entry *QueuedMessage + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + entries, err := s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + return err + } + for i := range entries { + if entries[i].ID == entryID { + found := entries[i] + entry = &found + return nil + } + } + return ErrEntryNotFound + }) + return entry, err +} + +type entryByIDRepository interface { + FindByID(context.Context, string) (*QueuedMessage, error) +} + +// FindEntryByID locates a queue entry after a session transfer has moved it +// out of its original session. +func (s *Service) FindEntryByID(ctx context.Context, entryID string) (*QueuedMessage, error) { + repository, ok := s.repo.(entryByIDRepository) + if !ok { + return nil, ErrEntryNotFound + } + return repository.FindByID(ctx, entryID) +} + +// ReferencedQueueAttachmentIDs returns attachment IDs referenced by pending +// entries in a session, excluding one entry. It is admission-aware so callers +// can use it while already holding the session admission lock. +func (s *Service) ReferencedQueueAttachmentIDs( + ctx context.Context, sessionID, excludedEntryID string, attachmentIDs []string, +) (map[string]struct{}, error) { + referenced := make(map[string]struct{}, len(attachmentIDs)) + if len(attachmentIDs) == 0 { + return referenced, nil + } + read := func(admittedCtx context.Context) error { + entries, err := s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + return err + } + for _, entry := range entries { + if entry.ID == excludedEntryID { + continue + } + for _, attachment := range entry.Attachments { + if attachment.AttachmentID != "" { + referenced[attachment.AttachmentID] = struct{}{} + } + } + } + return nil + } + if err := s.WithSessionAdmission(ctx, sessionID, read); err != nil { return nil, err } - for i := range entries { - if entries[i].ID == entryID { - entry := entries[i] - return &entry, nil - } - } - return nil, ErrEntryNotFound + return referenced, nil } // ClaimSendNow atomically claims the exact pending source snapshot for an @@ -1425,8 +2440,22 @@ func (s *Service) GetEntry(ctx context.Context, sessionID, entryID string) (*Que func (s *Service) ClaimSendNow(ctx context.Context, sessionID string, expected []QueuedMessage) (*SendNowClaim, error) { var claim *SendNowClaim err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + for _, entry := range expected { + if s.editLeaseBlocksEntryLocked(sessionID, entry.ID) { + return ErrEditConflict + } + } var err error claim, err = s.repo.ClaimSendNow(admittedCtx, sessionID, expected) + if err == nil && claim != nil { + s.editLeaseMu.Lock() + for _, source := range claim.Sources { + if !source.IsDurableLifecycle() { + s.deleteEditStateLocked(s.editLeaseKey(sessionID, source.ID)) + } + } + s.editLeaseMu.Unlock() + } return err }) if err != nil { @@ -1458,15 +2487,41 @@ func (s *Service) ClaimSendNowForSession(ctx context.Context, identity QueueSess // RestoreSendNowClaim restores every source from an interrupted replacement // dispatch. Durable lifecycle reservations are cleared as part of the same // repository operation. +func sendNowClaimSessionID(claim *SendNowClaim) (string, error) { + if claim == nil || len(claim.Sources) == 0 { + return "", ErrSendNowEmpty + } + sessionID := claim.Dispatch.SessionID + for _, source := range claim.Sources { + if source.SessionID == "" { + return "", ErrSendNowClaimChanged + } + if sessionID == "" { + sessionID = source.SessionID + } + if source.SessionID != sessionID { + return "", ErrSendNowClaimChanged + } + } + return sessionID, nil +} + func (s *Service) RestoreSendNowClaim(ctx context.Context, claim *SendNowClaim) error { - if err := s.repo.RestoreSendNowClaim(ctx, claim); err != nil { + if claim == nil { + return s.repo.RestoreSendNowClaim(ctx, claim) + } + sessionID, err := sendNowClaimSessionID(claim) + if err != nil { return err } - if claim != nil { - s.logger.Info("restored send-now queue claim", - zap.String("session_id", claim.Dispatch.SessionID), - zap.Int("source_count", len(claim.Sources))) + if err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + return s.repo.RestoreSendNowClaim(admittedCtx, claim) + }); err != nil { + return err } + s.logger.Info("restored send-now queue claim", + zap.String("session_id", sessionID), + zap.Int("source_count", len(claim.Sources))) return nil } @@ -1474,17 +2529,129 @@ func (s *Service) RestoreSendNowClaim(ctx context.Context, claim *SendNowClaim) // replacement prompt has been accepted. Ordinary sources were deleted at // claim time and therefore need no second acknowledgement. func (s *Service) AcknowledgeSendNowClaim(ctx context.Context, claim *SendNowClaim) error { - if err := s.repo.AcknowledgeSendNowClaim(ctx, claim); err != nil { + if claim == nil { + return s.repo.AcknowledgeSendNowClaim(ctx, claim) + } + sessionID, err := sendNowClaimSessionID(claim) + if err != nil { return err } - if claim != nil { - s.logger.Info("acknowledged send-now queue claim", - zap.String("session_id", claim.Dispatch.SessionID), - zap.Int("source_count", len(claim.Sources))) + if err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + return s.repo.AcknowledgeSendNowClaim(admittedCtx, claim) + }); err != nil { + return err } + s.logger.Info("acknowledged send-now queue claim", + zap.String("session_id", sessionID), + zap.Int("source_count", len(claim.Sources))) return nil } +type pendingSendNowClaimRepository interface { + ListPendingSendNowClaims(context.Context) ([]PendingSendNowClaim, error) + MarkPendingSendNowClaimAccepted(context.Context, *SendNowClaim) error + DeletePendingSendNowClaim(context.Context, *SendNowClaim) error +} + +// PendingSendNowClaimPersistenceAvailable reports whether ordinary claimed +// prompts can be recovered after the owning process exits. +func (s *Service) PendingSendNowClaimPersistenceAvailable() bool { + _, ok := s.repo.(pendingSendNowClaimRepository) + return ok +} + +// ListPendingSendNowClaims reloads interrupted Send Now claims after restart. +func (s *Service) ListPendingSendNowClaims(ctx context.Context) ([]PendingSendNowClaim, error) { + repo, ok := s.repo.(pendingSendNowClaimRepository) + if !ok { + return nil, errors.New("pending Send Now claim persistence unavailable") + } + return repo.ListPendingSendNowClaims(ctx) +} + +// MarkPendingSendNowClaimAccepted records the executor acceptance boundary +// before the potentially long-running prompt call returns. +func (s *Service) MarkPendingSendNowClaimAccepted(ctx context.Context, claim *SendNowClaim) error { + repo, ok := s.repo.(pendingSendNowClaimRepository) + if !ok { + return errors.New("pending Send Now claim persistence unavailable") + } + sessionID, err := sendNowClaimSessionID(claim) + if err != nil { + return err + } + return s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + return repo.MarkPendingSendNowClaimAccepted(admittedCtx, claim) + }) +} + +// DeletePendingSendNowClaim discards only the exact recovery record that can +// no longer be applied because a destructive queue generation change +// superseded it. +func (s *Service) DeletePendingSendNowClaim(ctx context.Context, claim *SendNowClaim) error { + repo, ok := s.repo.(pendingSendNowClaimRepository) + if !ok { + return errors.New("pending Send Now claim persistence unavailable") + } + sessionID, err := sendNowClaimSessionID(claim) + if err != nil { + return err + } + return s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + return repo.DeletePendingSendNowClaim(admittedCtx, claim) + }) +} + +type pendingQueueDispatchRepository interface { + ListPendingQueueDispatches(context.Context) ([]PendingQueueDispatch, error) + MarkPendingQueueDispatchAccepted(context.Context, *QueuedMessage) error + DeletePendingQueueDispatch(context.Context, *QueuedMessage) error +} + +// PendingQueueDispatchPersistenceAvailable reports whether ordinary dequeues +// survive a process exit before executor acceptance. +func (s *Service) PendingQueueDispatchPersistenceAvailable() bool { + _, ok := s.repo.(pendingQueueDispatchRepository) + return ok +} + +func (s *Service) ListPendingQueueDispatches(ctx context.Context) ([]PendingQueueDispatch, error) { + repo, ok := s.repo.(pendingQueueDispatchRepository) + if !ok { + return nil, errors.New("pending queue dispatch persistence unavailable") + } + return repo.ListPendingQueueDispatches(ctx) +} + +func (s *Service) MarkPendingQueueDispatchAccepted( + ctx context.Context, + msg *QueuedMessage, +) error { + repo, ok := s.repo.(pendingQueueDispatchRepository) + if !ok { + return nil + } + return s.WithSessionAdmission(ctx, msg.SessionID, func(admittedCtx context.Context) error { + return repo.MarkPendingQueueDispatchAccepted(admittedCtx, msg) + }) +} + +// DeletePendingQueueDispatch acknowledges the exact recovered or accepted +// ordinary dispatch attempt carried by msg. +func (s *Service) DeletePendingQueueDispatch(ctx context.Context, msg *QueuedMessage) error { + return s.WithSessionAdmission(ctx, msg.SessionID, func(admittedCtx context.Context) error { + return s.deletePendingQueueDispatch(admittedCtx, msg) + }) +} + +func (s *Service) deletePendingQueueDispatch(ctx context.Context, msg *QueuedMessage) error { + repo, ok := s.repo.(pendingQueueDispatchRepository) + if !ok || msg.IsDurableLifecycle() { + return nil + } + return repo.DeletePendingQueueDispatch(ctx, msg) +} + // GetEntryForSession returns an entry from an identity-bound repository // snapshot, preventing edits from observing a reincarnated session. func (s *Service) GetEntryForSession(ctx context.Context, identity QueueSessionIdentity, entryID string) (*QueuedMessage, error) { @@ -1505,6 +2672,10 @@ func (s *Service) GetEntryForSession(ctx context.Context, identity QueueSessionI // metadata replacements while retaining unrelated metadata keys. func (s *Service) UpdateMessageWithMetadata(ctx context.Context, sessionID, entryID, content string, attachments []MessageAttachment, metadataUpdates map[string]interface{}, queuedBy string) error { if err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + if s.editLeaseBlocksEntryLocked(sessionID, entryID) { + return ErrEditConflict + } + return s.repo.UpdateContentAndMetadata(admittedCtx, sessionID, entryID, content, attachments, metadataUpdates, queuedBy) }); err != nil { return err @@ -1515,6 +2686,39 @@ func (s *Service) UpdateMessageWithMetadata(ctx context.Context, sessionID, entr return nil } +// RemoveEntryWithEntry deletes a single entry and returns the exact snapshot +// removed under the same session admission lock. Callers that release +// entry-owned resources must use this method instead of GetEntry followed by +// RemoveEntry, because an edit between those calls can change attachments. +func (s *Service) RemoveEntryWithEntry(ctx context.Context, sessionID, entryID string) (*QueuedMessage, error) { + var removed *QueuedMessage + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + entries, err := s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + return err + } + for i := range entries { + if entries[i].ID != entryID { + continue + } + entry := entries[i] + if err := s.repo.DeleteByID(admittedCtx, sessionID, entryID); err != nil { + return err + } + removed = &entry + s.editLeaseMu.Lock() + s.deleteEditStateLocked(s.editLeaseKey(sessionID, entryID)) + s.editLeaseMu.Unlock() + return nil + } + return ErrEntryNotFound + }) + if err != nil { + return nil, err + } + return removed, nil +} + // UpdateMessageWithMetadataForSession edits content only for the exact session incarnation. func (s *Service) UpdateMessageWithMetadataForSession(ctx context.Context, identity QueueSessionIdentity, entryID, content string, attachments []MessageAttachment, metadataUpdates map[string]interface{}, queuedBy string) error { if err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error { @@ -1550,9 +2754,8 @@ func (s *Service) UpdateMessageWithMetadataForSessionWithClaim(ctx context.Conte // the rationale on the Repository.DeleteByID contract for why. Returns // ErrEntryNotFound when no entry matches. func (s *Service) RemoveEntry(ctx context.Context, sessionID, entryID string) error { - if err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { - return s.repo.DeleteByID(admittedCtx, sessionID, entryID) - }); err != nil { + _, err := s.RemoveEntryWithEntry(ctx, sessionID, entryID) + if err != nil { return err } s.logger.Info("queued entry removed", @@ -1588,9 +2791,19 @@ func (s *Service) MergeIntoAbove(ctx context.Context, sessionID, entryID, queued } var merged *QueuedMessage err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { - var err error - merged, err = s.repo.MergeIntoAbove(admittedCtx, sessionID, entryID, queuedBy) - return err + blocked, err := s.editLeaseBlocksMerge(admittedCtx, sessionID, entryID) + if err != nil { + return err + } + if blocked { + return ErrEditConflict + } + var mergeErr error + merged, mergeErr = s.repo.MergeIntoAbove(admittedCtx, sessionID, entryID, queuedBy) + if mergeErr == nil { + s.invalidateEditLease(sessionID, entryID) + } + return mergeErr }) if err != nil { return nil, err @@ -1631,6 +2844,9 @@ func (s *Service) ReorderEntries(ctx context.Context, sessionID string, orderedI return err } if err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + if s.editLeaseBlocksReorderLocked(sessionID) { + return ErrEditConflict + } return s.repo.ReorderEntries(admittedCtx, sessionID, orderedIDs) }); err != nil { return err @@ -1664,6 +2880,9 @@ func (s *Service) CancelAll(ctx context.Context, sessionID string) (int, error) err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { var err error n, err = s.repo.DeleteAllBySession(admittedCtx, sessionID) + if err == nil { + s.invalidateEditLeasesLocked(sessionID) + } return err }) if err != nil { @@ -1675,6 +2894,54 @@ func (s *Service) CancelAll(ctx context.Context, sessionID string) (int, error) return n, nil } +// PurgeSession removes all queue rows for a deleted session, including +// reserved lifecycle deliveries. It is intentionally distinct from +// CancelAll, which preserves in-flight durable rows for executor recovery. +func (s *Service) PurgeSession(ctx context.Context, sessionID string) (int, error) { + var removed int + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + var err error + removed, err = s.repo.PurgeSession(admittedCtx, sessionID) + if err == nil { + s.invalidateEditLeasesLocked(sessionID) + } + return err + }) + if err != nil { + return 0, err + } + return removed, nil +} + +// CancelAllWithEntries clears a session queue and returns the deleted rows so +// callers can release resources owned by those entries. +func (s *Service) CancelAllWithEntries(ctx context.Context, sessionID string) ([]QueuedMessage, error) { + var entries []QueuedMessage + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + var err error + entries, err = s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + return err + } + _, err = s.repo.DeleteAllBySession(admittedCtx, sessionID) + if err == nil { + deleted := entries[:0] + for i := range entries { + if !entries[i].IsReservedInFlight() { + deleted = append(deleted, entries[i]) + } + } + entries = deleted + s.invalidateEditLeasesLocked(sessionID) + } + return err + }) + if err != nil { + return nil, err + } + return entries, nil +} + // CancelAllForSession clears pending entries only for the exact session // incarnation and returns both removed and retained rows atomically. func (s *Service) CancelAllForSession(ctx context.Context, identity QueueSessionIdentity) (*QueueRemovalResult, error) { @@ -1709,37 +2976,58 @@ func (s *Service) PurgeDeletedSession(ctx context.Context, identity QueueSession func (s *Service) GetStatus(ctx context.Context, sessionID string) *QueueStatus { maxPerSession := s.MaxPerSession() mergeEnabled := s.MergeEnabled() - autoRun, autoRunErr := s.repo.GetAutoRun(ctx, sessionID) - if autoRunErr != nil { - s.logger.Error("get queue auto-run failed", - zap.String("session_id", sessionID), - zap.Error(autoRunErr)) - // Preserve pre-policy behavior if status storage is temporarily unreadable. - autoRun = true - } - entries, err := s.repo.ListBySession(ctx, sessionID) - if err != nil { - s.logger.Error("list queued failed", - zap.String("session_id", sessionID), - zap.Error(err)) - return &QueueStatus{Entries: []QueuedMessage{}, Count: 0, Max: maxPerSession, AutoRun: autoRun, MergeEnabled: mergeEnabled} - } - pending := make([]QueuedMessage, 0, len(entries)) - for _, entry := range entries { - // A reserved lifecycle row is already being delivered; listing it next - // to the message it produced reads as a stuck duplicate. - if entry.IsReservedInFlight() { - continue - } - pending = append(pending, entry) - } - return &QueueStatus{ - Entries: pending, - Count: len(pending), + status := &QueueStatus{ + Entries: []QueuedMessage{}, + Count: 0, Max: maxPerSession, - AutoRun: autoRun, + AutoRun: true, MergeEnabled: mergeEnabled, } + _ = s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + autoRun, autoRunErr := s.repo.GetAutoRun(admittedCtx, sessionID) + if autoRunErr != nil { + s.logger.Error("get queue auto-run failed", + zap.String("session_id", sessionID), + zap.Error(autoRunErr)) + // Preserve pre-policy behavior if status storage is temporarily unreadable. + autoRun = true + } + entries, err := s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + s.logger.Error("list queued failed", + zap.String("session_id", sessionID), + zap.Error(err)) + status = &QueueStatus{ + Entries: []QueuedMessage{}, + Count: 0, + Max: maxPerSession, + AutoRun: autoRun, + MergeEnabled: mergeEnabled, + } + return nil + } + pending := make([]QueuedMessage, 0, len(entries)) + for _, entry := range entries { + // A reserved lifecycle row is already being delivered; listing it next + // to the message it produced reads as a stuck duplicate. + if entry.IsReservedInFlight() { + continue + } + // Repositories may return shallow copies. The status is also handed to + // asynchronous event publishers, so detach maps and slices before the + // admission lock is released. + pending = append(pending, *cloneQueuedMessage(&entry)) + } + status = &QueueStatus{ + Entries: pending, + Count: len(pending), + Max: maxPerSession, + AutoRun: autoRun, + MergeEnabled: mergeEnabled, + } + return nil + }) + return status } // Snapshot returns an ordered status bound to one immutable session identity. @@ -1822,15 +3110,46 @@ func (s *Service) CountPendingByTask(ctx context.Context, taskID string) (int, e // Unlike GetStatus, it includes durable lifecycle rows reserved by an in-flight // delivery so a later restore cannot erase them before acknowledgement. func (s *Service) SnapshotSession(ctx context.Context, sessionID string) ([]QueuedMessage, *PendingMove, error) { - entries, err := s.repo.ListBySession(ctx, sessionID) - if err != nil { - return nil, nil, fmt.Errorf("snapshot queued messages: %w", err) + var entries []QueuedMessage + var move *PendingMove + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + rawEntries, err := s.repo.ListBySession(admittedCtx, sessionID) + if err != nil { + return fmt.Errorf("snapshot queued messages: %w", err) + } + entries = make([]QueuedMessage, 0, len(rawEntries)) + for i := range rawEntries { + entries = append(entries, *cloneQueuedMessage(&rawEntries[i])) + } + move, err = s.repo.GetPendingMove(admittedCtx, sessionID) + if err != nil { + return fmt.Errorf("snapshot pending move: %w", err) + } + if move != nil { + moveCopy := *move + move = &moveCopy + } + return nil + }) + return entries, move, err +} + +type ownedSessionTransferRepository interface { + transferSessionOwned(context.Context, string, string, string) error +} + +func (s *Service) transferRepositorySession( + ctx context.Context, + oldSessionID, newSessionID, operationID string, +) error { + if operationID == "" { + return s.repo.TransferSession(ctx, oldSessionID, newSessionID) } - move, err := s.repo.GetPendingMove(ctx, sessionID) - if err != nil { - return nil, nil, fmt.Errorf("snapshot pending move: %w", err) + repo, ok := s.repo.(ownedSessionTransferRepository) + if !ok { + return errors.New("owned session transfer unavailable") } - return entries, move, nil + return repo.transferSessionOwned(ctx, oldSessionID, newSessionID, operationID) } // SnapshotSessionForIdentity captures queue and deferred-move state for one exact incarnation. @@ -1848,7 +3167,58 @@ func (s *Service) SnapshotSessionForIdentity(ctx context.Context, identity Queue // old session — a transfer that no-ops without a signal would let the workflow // step move forward while the queue sticks behind. func (s *Service) TransferSession(ctx context.Context, oldSessionID, newSessionID string) error { - if err := s.repo.TransferSession(ctx, oldSessionID, newSessionID); err != nil { + return s.transferSession(ctx, "", oldSessionID, newSessionID, "", nil, nil) +} + +// TransferSessionWithPreparation runs preparation while both session +// admissions are held, then transfers queued state. If the queue transfer +// fails, rollback runs under the same admissions before they are released. +// Callers use this to keep external state, such as attachment ownership, +// synchronized with the queue move. +func (s *Service) TransferSessionWithPreparation( + ctx context.Context, + oldSessionID, newSessionID string, + prepare func(context.Context) error, + rollback func(context.Context) error, +) error { + return s.transferSession(ctx, "", oldSessionID, newSessionID, "", prepare, rollback) +} + +func (s *Service) transferSession( + ctx context.Context, + taskID string, + oldSessionID, newSessionID, operationID string, + prepare func(context.Context) error, + rollback func(context.Context) error, +) error { + err := s.withSessionAdmissions(ctx, oldSessionID, newSessionID, func(admittedCtx context.Context) error { + rollbackPreparation := func(transferErr error) error { + if rollback == nil { + return transferErr + } + rollbackErr := rollback(context.WithoutCancel(admittedCtx)) + if rollbackErr != nil { + s.logger.Warn("failed to roll back transfer preparation", + zap.String("from_session_id", oldSessionID), + zap.String("to_session_id", newSessionID), + zap.Error(rollbackErr)) + } + return errors.Join(transferErr, rollbackErr) + } + if prepare != nil { + if err := prepare(admittedCtx); err != nil { + return rollbackPreparation(err) + } + } + if err := s.transferRepositorySessionForTask( + admittedCtx, taskID, oldSessionID, newSessionID, operationID, + ); err != nil { + return rollbackPreparation(err) + } + s.invalidateEditLeasesLocked(oldSessionID, newSessionID) + return nil + }) + if err != nil { s.logger.Error("transfer session failed", zap.String("from_session_id", oldSessionID), zap.String("to_session_id", newSessionID), @@ -1861,6 +3231,276 @@ func (s *Service) TransferSession(ctx context.Context, oldSessionID, newSessionI return nil } +func (s *Service) transferRepositorySessionForTask( + ctx context.Context, + taskID, oldSessionID, newSessionID, operationID string, +) error { + // Transfers that do not have a durable compensation record can use the + // immutable session identities. This fences a workflow handoff to one task + // and lets repositories reject a stale or cross-task destination. Durable + // transfers keep the owned-operation path because it also carries the + // compensation lease and operation token. + if taskID != "" && operationID == "" { + source, sourceErr := s.repo.ResolveSessionIdentity(ctx, taskID, oldSessionID) + if sourceErr != nil && !errors.Is(sourceErr, ErrSessionIdentityMismatch) { + return sourceErr + } + destination, destinationErr := s.repo.ResolveSessionIdentity(ctx, taskID, newSessionID) + if destinationErr != nil && !errors.Is(destinationErr, ErrSessionIdentityMismatch) { + return destinationErr + } + if sourceErr == nil && destinationErr == nil { + return s.repo.TransferSessionIdentities(ctx, source, destination) + } + } + return s.transferRepositorySession(ctx, oldSessionID, newSessionID, operationID) +} + +type attachmentCleanupRepository interface { + UpsertAttachmentCleanup(context.Context, AttachmentCleanup) error + DeleteAttachmentCleanup(context.Context, string, string, string) error + ListAttachmentCleanups(context.Context) ([]AttachmentCleanup, error) +} + +// AttachmentCleanupPersistenceAvailable reports whether this service's +// repository can preserve cleanup work across handler restarts. +func (s *Service) AttachmentCleanupPersistenceAvailable() bool { + _, ok := s.repo.(attachmentCleanupRepository) + return ok +} + +// UpsertAttachmentCleanup persists a cleanup obligation before its retry +// worker is allowed to outlive the current handler. +func (s *Service) UpsertAttachmentCleanup(ctx context.Context, cleanup AttachmentCleanup) error { + repo, ok := s.repo.(attachmentCleanupRepository) + if !ok { + return errors.New("attachment cleanup persistence unavailable") + } + return repo.UpsertAttachmentCleanup(ctx, cleanup) +} + +// DeleteAttachmentCleanup acknowledges one completed cleanup obligation. +func (s *Service) DeleteAttachmentCleanup( + ctx context.Context, + sessionID, entryID, operationID string, +) error { + repo, ok := s.repo.(attachmentCleanupRepository) + if !ok { + return errors.New("attachment cleanup persistence unavailable") + } + return repo.DeleteAttachmentCleanup(ctx, sessionID, entryID, operationID) +} + +// ListAttachmentCleanups reloads cleanup obligations after process restart. +func (s *Service) ListAttachmentCleanups(ctx context.Context) ([]AttachmentCleanup, error) { + repo, ok := s.repo.(attachmentCleanupRepository) + if !ok { + return nil, errors.New("attachment cleanup persistence unavailable") + } + return repo.ListAttachmentCleanups(ctx) +} + +type attachmentCleanupLocatorRepository interface { + GetAttachmentCleanup(context.Context, string, string, string) (*AttachmentCleanup, error) +} + +// GetAttachmentCleanup reloads one cleanup by its immutable acknowledgement key. +func (s *Service) GetAttachmentCleanup( + ctx context.Context, + sessionID, entryID, operationID string, +) (*AttachmentCleanup, error) { + repo, ok := s.repo.(attachmentCleanupLocatorRepository) + if !ok { + return nil, errors.New("attachment cleanup persistence unavailable") + } + return repo.GetAttachmentCleanup(ctx, sessionID, entryID, operationID) +} + +type sessionTransferCompensationRepository interface { + UpsertSessionTransferCompensation(context.Context, SessionTransferCompensation) error + DeleteSessionTransferCompensation(context.Context, string, string, string, string) error + ListSessionTransferCompensations(context.Context) ([]SessionTransferCompensation, error) +} + +type sessionTransferCompensationRecoveryRepository interface { + claimSessionTransferCompensationRecovery(context.Context, SessionTransferCompensation) (string, error) + renewSessionTransferCompensationLease(context.Context, SessionTransferCompensation, string) error + deleteSessionTransferCompensationWithOwner(context.Context, string, string, string, string, string) error +} + +// SessionTransferCompensationPersistenceAvailable reports whether interrupted +// queue and attachment transfers can be reconciled after process restart. +func (s *Service) SessionTransferCompensationPersistenceAvailable() bool { + _, ok := s.repo.(sessionTransferCompensationRepository) + return ok +} + +// UpsertSessionTransferCompensation persists intent before attachment ownership +// changes outside the queue repository transaction. +func (s *Service) UpsertSessionTransferCompensation( + ctx context.Context, + compensation SessionTransferCompensation, +) error { + repo, ok := s.repo.(sessionTransferCompensationRepository) + if !ok { + return errors.New("session transfer compensation persistence unavailable") + } + return repo.UpsertSessionTransferCompensation(ctx, compensation) +} + +// DeleteSessionTransferCompensation acknowledges a reconciled transfer. +func (s *Service) DeleteSessionTransferCompensation( + ctx context.Context, + operationID, taskID, fromSessionID, toSessionID string, +) error { + repo, ok := s.repo.(sessionTransferCompensationRepository) + if !ok { + return errors.New("session transfer compensation persistence unavailable") + } + return repo.DeleteSessionTransferCompensation( + ctx, operationID, taskID, fromSessionID, toSessionID, + ) +} + +// RenewSessionTransferCompensationLease keeps external transfer work owned +// until its queue mutation or recovery acknowledgement completes. +func (s *Service) RenewSessionTransferCompensationLease( + ctx context.Context, + compensation SessionTransferCompensation, + ownerID string, +) error { + repo, ok := s.repo.(sessionTransferCompensationRecoveryRepository) + if !ok { + return errors.New("session transfer compensation recovery unavailable") + } + return repo.renewSessionTransferCompensationLease(ctx, compensation, ownerID) +} + +// ClaimSessionTransferCompensationRecovery claims an expired transfer +// compensation before its external attachment recovery starts. +func (s *Service) ClaimSessionTransferCompensationRecovery( + ctx context.Context, + compensation SessionTransferCompensation, +) (string, error) { + repo, ok := s.repo.(sessionTransferCompensationRecoveryRepository) + if !ok { + return "", errors.New("session transfer compensation recovery unavailable") + } + return repo.claimSessionTransferCompensationRecovery(ctx, compensation) +} + +// DeleteClaimedSessionTransferCompensation acknowledges recovered attachment +// state only while the caller owns the compensation recovery lease. + +// MaintainSessionTransferCompensationLease renews an owner-scoped lease until +// the caller completes its external attachment operation. Renewal failure +// cancels the returned context and is returned by the stop function. +func (s *Service) MaintainSessionTransferCompensationLease( + ctx context.Context, + compensation SessionTransferCompensation, + ownerID string, +) (context.Context, func() error) { + return s.maintainSessionTransferCompensationLease( + ctx, compensation, ownerID, nil, sessionTransferCompensationLeaseDuration/2, + ) +} + +// MaintainSessionTransferCompensationLeaseWithCancel is the recovery variant +// that also cancels the owning operation when lease renewal fails. +func (s *Service) MaintainSessionTransferCompensationLeaseWithCancel( + ctx context.Context, + compensation SessionTransferCompensation, + ownerID string, + cancelOwner context.CancelFunc, +) (context.Context, func() error) { + return s.maintainSessionTransferCompensationLease( + ctx, compensation, ownerID, cancelOwner, sessionTransferCompensationLeaseDuration/2, + ) +} + +func (s *Service) maintainSessionTransferCompensationLease( + ctx context.Context, + compensation SessionTransferCompensation, + ownerID string, + cancelOwner context.CancelFunc, + renewInterval time.Duration, +) (context.Context, func() error) { + operationCtx, cancel := context.WithCancel(ctx) + stop := make(chan struct{}) + done := make(chan error, 1) + var stopOnce sync.Once + var stopErr error + go func() { + ticker := time.NewTicker(renewInterval) + defer ticker.Stop() + var err error + for { + select { + case <-stop: + done <- nil + return + case <-operationCtx.Done(): + done <- nil + return + case <-ticker.C: + renewCtx, renewCancel := context.WithTimeout( + operationCtx, sessionTransferCompensationLeaseDuration/3, + ) + err = s.RenewSessionTransferCompensationLease( + renewCtx, compensation, ownerID, + ) + renewCancel() + if err != nil { + cancel() + if cancelOwner != nil { + cancelOwner() + } + done <- err + return + } + } + } + }() + return operationCtx, func() error { + stopOnce.Do(func() { + close(stop) + stopErr = <-done + cancel() + }) + return stopErr + } +} + +func (s *Service) DeleteClaimedSessionTransferCompensation( + ctx context.Context, + compensation SessionTransferCompensation, + ownerID string, +) error { + repo, ok := s.repo.(sessionTransferCompensationRecoveryRepository) + if !ok { + return errors.New("session transfer compensation recovery unavailable") + } + return repo.deleteSessionTransferCompensationWithOwner( + ctx, + compensation.OperationID, + ownerID, + compensation.TaskID, + compensation.FromSessionID, + compensation.ToSessionID, + ) +} + +// ListSessionTransferCompensations reloads interrupted transfers after restart. +func (s *Service) ListSessionTransferCompensations( + ctx context.Context, +) ([]SessionTransferCompensation, error) { + repo, ok := s.repo.(sessionTransferCompensationRepository) + if !ok { + return nil, errors.New("session transfer compensation persistence unavailable") + } + return repo.ListSessionTransferCompensations(ctx) +} + func (s *Service) validateSessionIdentity(ctx context.Context, identity QueueSessionIdentity) error { current, err := s.repo.ResolveSessionIdentity(ctx, identity.TaskID, identity.SessionID) if err != nil { @@ -1899,7 +3539,14 @@ func (s *Service) TransferSessionIdentities(ctx context.Context, source, destina // RestoreSession replaces a session's queue and pending move from a snapshot, // preserving queued-message identity fields. func (s *Service) RestoreSession(ctx context.Context, sessionID string, entries []QueuedMessage, pendingMove *PendingMove) error { - if err := s.repo.ReplaceSession(ctx, sessionID, entries, pendingMove); err != nil { + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + if err := s.repo.ReplaceSession(admittedCtx, sessionID, entries, pendingMove); err != nil { + return err + } + s.invalidateEditLeasesLocked(sessionID) + return nil + }) + if err != nil { s.logger.Error("restore session queue failed", zap.String("session_id", sessionID), zap.Error(err)) @@ -1933,7 +3580,9 @@ func (s *Service) RestoreSessionForIdentity(ctx context.Context, identity QueueS // SetPendingMove records a pending move for a session (replaces any existing one). // The move is applied by handleAgentReady when the agent's current turn completes. func (s *Service) SetPendingMove(ctx context.Context, sessionID string, move *PendingMove) error { - if err := s.repo.SetPendingMove(ctx, sessionID, move); err != nil { + if err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + return s.repo.SetPendingMove(admittedCtx, sessionID, move) + }); err != nil { s.logger.Error("set pending move failed", zap.String("session_id", sessionID), zap.Error(err)) @@ -1950,7 +3599,12 @@ func (s *Service) SetPendingMove(ctx context.Context, sessionID string, move *Pe // removing it. It distinguishes an empty queue from a storage failure so // callers that must preserve deferred workflow state can fail closed. func (s *Service) GetPendingMoveWithError(ctx context.Context, sessionID string) (*PendingMove, bool, error) { - move, err := s.repo.GetPendingMove(ctx, sessionID) + var move *PendingMove + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + var err error + move, err = s.repo.GetPendingMove(admittedCtx, sessionID) + return err + }) if err != nil { return nil, false, err } @@ -1974,7 +3628,12 @@ func (s *Service) GetPendingMove(ctx context.Context, sessionID string) (*Pendin // TakePendingMove retrieves and removes the pending move for a session. func (s *Service) TakePendingMove(ctx context.Context, sessionID string) (*PendingMove, bool) { - move, err := s.repo.TakePendingMove(ctx, sessionID) + var move *PendingMove + err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error { + var err error + move, err = s.repo.TakePendingMove(admittedCtx, sessionID) + return err + }) if err != nil { s.logger.Error("take pending move failed", zap.String("session_id", sessionID), diff --git a/apps/backend/internal/orchestrator/messagequeue/service_admission_regression_test.go b/apps/backend/internal/orchestrator/messagequeue/service_admission_regression_test.go new file mode 100644 index 0000000000..33e72afc5e --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/service_admission_regression_test.go @@ -0,0 +1,328 @@ +package messagequeue + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type admissionProbeRepository struct { + Repository + listCalled chan struct{} + restoreCalled chan struct{} + acknowledgeCalled chan struct{} + setPendingCalled chan struct{} + getPendingCalled chan struct{} + takePendingCalled chan struct{} +} + +func (r *admissionProbeRepository) ListBySession(ctx context.Context, sessionID string) ([]QueuedMessage, error) { + select { + case r.listCalled <- struct{}{}: + default: + } + return r.Repository.ListBySession(ctx, sessionID) +} + +func (r *admissionProbeRepository) RestoreSendNowClaim(ctx context.Context, claim *SendNowClaim) error { + select { + case r.restoreCalled <- struct{}{}: + default: + } + return r.Repository.RestoreSendNowClaim(ctx, claim) +} + +func (r *admissionProbeRepository) AcknowledgeSendNowClaim(ctx context.Context, claim *SendNowClaim) error { + select { + case r.acknowledgeCalled <- struct{}{}: + default: + } + return r.Repository.AcknowledgeSendNowClaim(ctx, claim) +} + +func (r *admissionProbeRepository) SetPendingMove(ctx context.Context, sessionID string, move *PendingMove) error { + select { + case r.setPendingCalled <- struct{}{}: + default: + } + return r.Repository.SetPendingMove(ctx, sessionID, move) +} + +func (r *admissionProbeRepository) GetPendingMove(ctx context.Context, sessionID string) (*PendingMove, error) { + select { + case r.getPendingCalled <- struct{}{}: + default: + } + return r.Repository.GetPendingMove(ctx, sessionID) +} + +func (r *admissionProbeRepository) TakePendingMove(ctx context.Context, sessionID string) (*PendingMove, error) { + select { + case r.takePendingCalled <- struct{}{}: + default: + } + return r.Repository.TakePendingMove(ctx, sessionID) +} + +func newAdmissionProbeRepository() *admissionProbeRepository { + return &admissionProbeRepository{ + Repository: NewMemoryRepository(), + listCalled: make(chan struct{}, 1), + restoreCalled: make(chan struct{}, 1), + acknowledgeCalled: make(chan struct{}, 1), + setPendingCalled: make(chan struct{}, 1), + getPendingCalled: make(chan struct{}, 1), + takePendingCalled: make(chan struct{}, 1), + } +} + +func TestServiceQueueOperationsWaitForAdmission(t *testing.T) { + tests := []struct { + name string + probe func(*admissionProbeRepository) <-chan struct{} + prepare func(context.Context, *Service, *admissionProbeRepository) func(context.Context) error + }{ + { + name: "get entry", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.listCalled }, + prepare: func(ctx context.Context, svc *Service, _ *admissionProbeRepository) func(context.Context) error { + entry, err := svc.QueueMessage(ctx, "session", "task", "content", "", QueuedByUser, false, nil) + require.NoError(t, err) + return func(operationCtx context.Context) error { + _, err := svc.GetEntry(operationCtx, "session", entry.ID) + return err + } + }, + }, + { + name: "restore send-now claim", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.restoreCalled }, + prepare: func(ctx context.Context, svc *Service, _ *admissionProbeRepository) func(context.Context) error { + entry, err := svc.QueueMessage(ctx, "session", "task", "content", "", QueuedByUser, false, nil) + require.NoError(t, err) + claim, err := svc.ClaimSendNow(ctx, "session", []QueuedMessage{*entry}) + require.NoError(t, err) + return func(operationCtx context.Context) error { return svc.RestoreSendNowClaim(operationCtx, claim) } + }, + }, + { + name: "acknowledge send-now claim", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.acknowledgeCalled }, + prepare: func(ctx context.Context, svc *Service, _ *admissionProbeRepository) func(context.Context) error { + entry, err := svc.QueueMessage(ctx, "session", "task", "content", "", QueuedByUser, false, nil) + require.NoError(t, err) + claim, err := svc.ClaimSendNow(ctx, "session", []QueuedMessage{*entry}) + require.NoError(t, err) + return func(operationCtx context.Context) error { return svc.AcknowledgeSendNowClaim(operationCtx, claim) } + }, + }, + { + name: "set pending move", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.setPendingCalled }, + prepare: func(_ context.Context, svc *Service, _ *admissionProbeRepository) func(context.Context) error { + return func(operationCtx context.Context) error { + return svc.SetPendingMove(operationCtx, "session", &PendingMove{TaskID: "task"}) + } + }, + }, + { + name: "get pending move", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.getPendingCalled }, + prepare: func(ctx context.Context, svc *Service, _ *admissionProbeRepository) func(context.Context) error { + require.NoError(t, svc.SetPendingMove(ctx, "session", &PendingMove{TaskID: "task"})) + return func(operationCtx context.Context) error { + _, _ = svc.GetPendingMove(operationCtx, "session") + return nil + } + }, + }, + { + name: "take pending move", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.takePendingCalled }, + prepare: func(ctx context.Context, svc *Service, _ *admissionProbeRepository) func(context.Context) error { + require.NoError(t, svc.SetPendingMove(ctx, "session", &PendingMove{TaskID: "task"})) + return func(operationCtx context.Context) error { + _, _ = svc.TakePendingMove(operationCtx, "session") + return nil + } + }, + }, + { + name: "get status", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.listCalled }, + prepare: func(_ context.Context, svc *Service, _ *admissionProbeRepository) func(context.Context) error { + return func(operationCtx context.Context) error { + _ = svc.GetStatus(operationCtx, "session") + return nil + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := newAdmissionProbeRepository() + svc := newAutoMergeTestServiceWithRepository(t, repo, 10) + svc.SetAutoMergeEnabled(false) + ctx := context.Background() + op := tt.prepare(ctx, svc, repo) + + holderStarted := make(chan struct{}) + releaseHolder := make(chan struct{}) + holderDone := make(chan error, 1) + go func() { + holderDone <- svc.WithSessionAdmission(ctx, "session", func(context.Context) error { + close(holderStarted) + <-releaseHolder + return nil + }) + }() + <-holderStarted + + opDone := make(chan error, 1) + go func() { opDone <- op(ctx) }() + select { + case <-tt.probe(repo): + t.Fatal("queue operation reached repository while admission was held") + case <-time.After(100 * time.Millisecond): + } + + close(releaseHolder) + require.NoError(t, <-holderDone) + require.NoError(t, <-opDone) + select { + case <-tt.probe(repo): + case <-time.After(time.Second): + t.Fatal("queue operation did not reach repository after admission was released") + } + }) + } +} +func TestRestoreSendNowClaimUsesSourceSessionAdmission(t *testing.T) { + repo := newAdmissionProbeRepository() + svc := newAutoMergeTestServiceWithRepository(t, repo, 10) + svc.SetAutoMergeEnabled(false) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session", "task", "content", "", QueuedByUser, false, nil) + require.NoError(t, err) + claim := &SendNowClaim{Sources: []QueuedMessage{*entry}} + + holderStarted := make(chan struct{}) + releaseHolder := make(chan struct{}) + holderDone := make(chan error, 1) + go func() { + holderDone <- svc.WithSessionAdmission(ctx, "session", func(context.Context) error { + close(holderStarted) + <-releaseHolder + return nil + }) + }() + <-holderStarted + + restoreDone := make(chan error, 1) + go func() { restoreDone <- svc.RestoreSendNowClaim(ctx, claim) }() + select { + case <-repo.restoreCalled: + t.Fatal("pending Send Now restore reached repository while session admission was held") + case <-time.After(100 * time.Millisecond): + } + + close(releaseHolder) + require.NoError(t, <-holderDone) + require.NoError(t, <-restoreDone) + select { + case <-repo.restoreCalled: + case <-time.After(time.Second): + t.Fatal("pending Send Now restore did not reach repository after admission was released") + } +} +func TestSendNowClaimIdentityValidationRejectsBeforeRepository(t *testing.T) { + tests := []struct { + name string + claim func(QueuedMessage) *SendNowClaim + }{ + { + name: "dispatch session mismatch", + claim: func(entry QueuedMessage) *SendNowClaim { + return &SendNowClaim{ + Dispatch: QueuedMessage{SessionID: "other-session"}, + Sources: []QueuedMessage{entry}, + } + }, + }, + { + name: "mixed source sessions", + claim: func(entry QueuedMessage) *SendNowClaim { + other := entry + other.SessionID = "other-session" + return &SendNowClaim{Sources: []QueuedMessage{entry, other}} + }, + }, + { + name: "empty source session with dispatch", + claim: func(entry QueuedMessage) *SendNowClaim { + entry.SessionID = "" + return &SendNowClaim{ + Dispatch: QueuedMessage{SessionID: "session"}, + Sources: []QueuedMessage{entry}, + } + }, + }, + { + name: "empty source session without dispatch", + claim: func(entry QueuedMessage) *SendNowClaim { + entry.SessionID = "" + return &SendNowClaim{Sources: []QueuedMessage{entry}} + }, + }, + } + operations := []struct { + name string + probe func(*admissionProbeRepository) <-chan struct{} + call func(context.Context, *Service, *SendNowClaim) error + }{ + { + name: "restore", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.restoreCalled }, + call: func(ctx context.Context, svc *Service, claim *SendNowClaim) error { + return svc.RestoreSendNowClaim(ctx, claim) + }, + }, + { + name: "acknowledge", + probe: func(repo *admissionProbeRepository) <-chan struct{} { return repo.acknowledgeCalled }, + call: func(ctx context.Context, svc *Service, claim *SendNowClaim) error { + return svc.AcknowledgeSendNowClaim(ctx, claim) + }, + }, + } + + for _, operation := range operations { + t.Run(operation.name, func(t *testing.T) { + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + repo := newAdmissionProbeRepository() + svc := newAutoMergeTestServiceWithRepository(t, repo, 10) + svc.SetAutoMergeEnabled(false) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "session", "task", "content", "", QueuedByUser, false, nil) + require.NoError(t, err) + before := svc.GetStatus(ctx, entry.SessionID) + + err = operation.call(ctx, svc, test.claim(*entry)) + require.ErrorIs(t, err, ErrSendNowClaimChanged) + select { + case <-operation.probe(repo): + t.Fatal("malformed Send Now claim reached repository") + default: + } + + after := svc.GetStatus(ctx, entry.SessionID) + require.Equal(t, before.Entries, after.Entries) + }) + } + }) + } +} diff --git a/apps/backend/internal/orchestrator/messagequeue/service_session_transfer.go b/apps/backend/internal/orchestrator/messagequeue/service_session_transfer.go new file mode 100644 index 0000000000..e82f2ea697 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/service_session_transfer.go @@ -0,0 +1,306 @@ +package messagequeue + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +type durableSessionTransferState struct { + operationID string + compensation *SessionTransferCompensation + attachmentIDs []string + preparationCalled bool + rollbackSucceeded bool + operationCtx context.Context + stopLeaseRenewal func() error + leaseRenewalErr error + cancelTransfer context.CancelFunc +} + +// TransferSessionWithDurablePreparation preserves the original callback +// contract for callers whose external state is not attachment-scoped. +func (s *Service) TransferSessionWithDurablePreparation( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + prepare func(context.Context) error, + rollback func(context.Context) error, +) error { + return s.TransferSessionWithDurableAttachmentPreparation( + ctx, + taskID, + oldSessionID, + newSessionID, + func(callbackCtx context.Context, _ []string) error { + if prepare == nil { + return nil + } + return prepare(callbackCtx) + }, + func(callbackCtx context.Context, _ []string) error { + if rollback == nil { + return nil + } + return rollback(callbackCtx) + }, + ) +} + +// TransferSessionWithDurableAttachmentPreparation keeps attachment ownership +// aligned with every queue row, including lifecycle rows hidden from UI status. +// Durable repositories record the exact attachment set before external state +// changes so rollback and startup recovery cannot move unrelated claims. +func (s *Service) TransferSessionWithDurableAttachmentPreparation( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + prepare func(context.Context, []string) error, + rollback func(context.Context, []string) error, +) error { + state := &durableSessionTransferState{} + transferCtx, cancelTransfer := context.WithCancel(ctx) + defer cancelTransfer() + state.cancelTransfer = cancelTransfer + if s.SessionTransferCompensationPersistenceAvailable() { + state.operationID = uuid.NewString() + } + err := s.transferSession( + transferCtx, + taskID, + oldSessionID, + newSessionID, + state.operationID, + func(admittedCtx context.Context) error { + return s.prepareDurableSessionTransfer( + admittedCtx, taskID, oldSessionID, newSessionID, prepare, state, + ) + }, + func(rollbackCtx context.Context) error { + if !state.preparationCalled || rollback == nil { + state.rollbackSucceeded = true + return nil + } + rollbackErr := rollback(rollbackCtx, state.attachmentIDs) + state.rollbackSucceeded = rollbackErr == nil + return rollbackErr + }, + ) + if state.stopLeaseRenewal != nil { + state.leaseRenewalErr = state.stopLeaseRenewal() + state.stopLeaseRenewal = nil + if state.leaseRenewalErr != nil { + err = errors.Join(err, fmt.Errorf("session transfer lease lost: %w", state.leaseRenewalErr)) + } + } + if err != nil { + if state.compensation != nil && + state.rollbackSucceeded && + state.leaseRenewalErr == nil && + !errors.Is(err, ErrSessionTransferOwnershipLost) { + err = errors.Join(err, s.deleteSessionTransferCompensation(context.WithoutCancel(ctx), *state.compensation)) + } + return err + } + if state.compensation != nil { + if deleteErr := s.deleteSessionTransferCompensation(context.WithoutCancel(ctx), *state.compensation); deleteErr != nil { + s.logger.Warn("session transfer committed with pending compensation record", zap.Error(deleteErr)) + } + } + return nil +} + +func (s *Service) startSessionTransferCompensation( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + state *durableSessionTransferState, +) error { + if !s.SessionTransferCompensationPersistenceAvailable() { + return nil + } + state.compensation = &SessionTransferCompensation{ + OperationID: state.operationID, TaskID: taskID, + FromSessionID: oldSessionID, ToSessionID: newSessionID, + CreatedAt: time.Now().UTC(), + } + if err := s.UpsertSessionTransferCompensation(ctx, *state.compensation); err != nil { + return fmt.Errorf("persist session transfer fence: %w", err) + } + renewInterval := sessionTransferCompensationLeaseDuration / 2 + if s.sessionTransferCompensationLeaseRenewInterval > 0 { + renewInterval = s.sessionTransferCompensationLeaseRenewInterval + } + state.operationCtx, state.stopLeaseRenewal = s.maintainSessionTransferCompensationLease( + ctx, *state.compensation, state.operationID, state.cancelTransfer, renewInterval, + ) + return nil +} + +func (s *Service) prepareDurableSessionTransfer( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + prepare func(context.Context, []string) error, + state *durableSessionTransferState, +) error { + if err := s.startSessionTransferCompensation( + ctx, taskID, oldSessionID, newSessionID, state, + ); err != nil { + return err + } + entries, err := s.repo.ListBySession(ctx, oldSessionID) + if err != nil { + return fmt.Errorf("snapshot session queue for transfer: %w", err) + } + entryIDs := make([]string, 0, len(entries)) + seenEntryIDs := make(map[string]struct{}, len(entries)) + attachmentIDs := make([]string, 0) + seenAttachmentIDs := make(map[string]struct{}) + for _, entry := range entries { + entryIDs = append(entryIDs, entry.ID) + seenEntryIDs[entry.ID] = struct{}{} + attachmentIDs = appendTransferAttachmentIDs(attachmentIDs, seenAttachmentIDs, entry.Attachments) + } + dispatchEntryIDs, dispatchAttachments, err := s.snapshotPendingDispatchTransfer( + ctx, oldSessionID, seenEntryIDs, seenAttachmentIDs, + ) + if err != nil { + return err + } + entryIDs = append(entryIDs, dispatchEntryIDs...) + attachmentIDs = append(attachmentIDs, dispatchAttachments...) + cleanupEntryIDs, cleanupAttachments, cleanupLocators, err := s.snapshotAttachmentCleanupTransfer( + ctx, taskID, oldSessionID, seenEntryIDs, seenAttachmentIDs, + ) + if err != nil { + return err + } + entryIDs = append(entryIDs, cleanupEntryIDs...) + attachmentIDs = append(attachmentIDs, cleanupAttachments...) + state.attachmentIDs = attachmentIDs + if state.compensation != nil { + state.compensation.EntryIDs = entryIDs + state.compensation.AttachmentIDs = attachmentIDs + state.compensation.CleanupLocators = cleanupLocators + if err := s.UpsertSessionTransferCompensation(ctx, *state.compensation); err != nil { + return fmt.Errorf("persist session transfer compensation: %w", err) + } + } + if prepare == nil { + return nil + } + state.preparationCalled = true + prepareCtx := ctx + if state.operationCtx != nil { + prepareCtx = state.operationCtx + } + if err := prepare(prepareCtx, attachmentIDs); err != nil { + return fmt.Errorf("prepare durable session transfer: %w", err) + } + if state.operationCtx != nil { + select { + case <-state.operationCtx.Done(): + return state.operationCtx.Err() + default: + } + } + return nil +} + +func (s *Service) snapshotPendingDispatchTransfer( + ctx context.Context, + sessionID string, + seenEntryIDs, seenAttachmentIDs map[string]struct{}, +) ([]string, []string, error) { + if !s.PendingQueueDispatchPersistenceAvailable() { + return nil, nil, nil + } + pending, err := s.ListPendingQueueDispatches(ctx) + if err != nil { + return nil, nil, fmt.Errorf("snapshot in-flight queue dispatches for transfer: %w", err) + } + entryIDs := make([]string, 0) + attachmentIDs := make([]string, 0) + for _, dispatch := range pending { + entry := dispatch.Message + if entry.SessionID != sessionID { + continue + } + if _, seen := seenEntryIDs[entry.ID]; !seen { + entryIDs = append(entryIDs, entry.ID) + seenEntryIDs[entry.ID] = struct{}{} + } + attachmentIDs = appendTransferAttachmentIDs(attachmentIDs, seenAttachmentIDs, entry.Attachments) + } + return entryIDs, attachmentIDs, nil +} + +func (s *Service) snapshotAttachmentCleanupTransfer( + ctx context.Context, + taskID, sessionID string, + seenEntryIDs, seenAttachmentIDs map[string]struct{}, +) ([]string, []string, []AttachmentCleanupLocator, error) { + if !s.AttachmentCleanupPersistenceAvailable() { + return nil, nil, nil, nil + } + cleanups, err := s.ListAttachmentCleanups(ctx) + if err != nil { + return nil, nil, nil, fmt.Errorf("snapshot attachment cleanups for transfer: %w", err) + } + entryIDs := make([]string, 0) + attachmentIDs := make([]string, 0) + locators := make([]AttachmentCleanupLocator, 0) + for _, cleanup := range cleanups { + currentSessionID := cleanup.CurrentSessionID + if currentSessionID == "" { + currentSessionID = cleanup.SessionID + } + if cleanup.TaskID != taskID || currentSessionID != sessionID { + continue + } + if _, seen := seenEntryIDs[cleanup.EntryID]; !seen { + entryIDs = append(entryIDs, cleanup.EntryID) + seenEntryIDs[cleanup.EntryID] = struct{}{} + } + attachmentIDs = appendTransferAttachmentIDs(attachmentIDs, seenAttachmentIDs, cleanup.Attachments) + locators = append(locators, AttachmentCleanupLocator{ + SessionID: cleanup.SessionID, EntryID: cleanup.EntryID, OperationID: cleanup.OperationID, + }) + } + return entryIDs, attachmentIDs, locators, nil +} + +func appendTransferAttachmentIDs( + ids []string, + seen map[string]struct{}, + attachments []MessageAttachment, +) []string { + for _, attachment := range attachments { + if attachment.AttachmentID == "" { + continue + } + if _, ok := seen[attachment.AttachmentID]; ok { + continue + } + seen[attachment.AttachmentID] = struct{}{} + ids = append(ids, attachment.AttachmentID) + } + return ids +} + +func (s *Service) deleteSessionTransferCompensation( + ctx context.Context, + compensation SessionTransferCompensation, +) error { + if err := s.DeleteSessionTransferCompensation( + ctx, + compensation.OperationID, + compensation.TaskID, + compensation.FromSessionID, + compensation.ToSessionID, + ); err != nil { + return fmt.Errorf("delete session transfer compensation: %w", err) + } + return nil +} diff --git a/apps/backend/internal/orchestrator/messagequeue/service_status_snapshot_test.go b/apps/backend/internal/orchestrator/messagequeue/service_status_snapshot_test.go new file mode 100644 index 0000000000..7b930ec6d4 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/service_status_snapshot_test.go @@ -0,0 +1,74 @@ +package messagequeue + +import ( + "context" + "testing" + + "github.com/kandev/kandev/internal/common/logger" + "github.com/stretchr/testify/require" +) + +func TestGetStatusReturnsDetachedEntrySnapshots(t *testing.T) { + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console", OutputPath: "stderr"}) + require.NoError(t, err) + svc := NewServiceMemory(log) + ctx := context.Background() + _, err = svc.QueueMessageWithMetadata(ctx, "session", "task", "content", "", QueuedByUser, false, + []MessageAttachment{{AttachmentID: "attachment-1", Name: "original.txt"}}, + map[string]interface{}{"nested": map[string]interface{}{"value": "original"}}, + ) + require.NoError(t, err) + + status := svc.GetStatus(ctx, "session") + require.Len(t, status.Entries, 1) + status.Entries[0].Attachments[0].Name = "mutated.txt" + status.Entries[0].Metadata["nested"].(map[string]interface{})["value"] = "mutated" + + fresh := svc.GetStatus(ctx, "session") + require.Equal(t, "original.txt", fresh.Entries[0].Attachments[0].Name) + require.Equal(t, "original", fresh.Entries[0].Metadata["nested"].(map[string]interface{})["value"]) +} +func TestGetStatusReturnsDetachedTypedMetadataSnapshots(t *testing.T) { + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console", OutputPath: "stderr"}) + require.NoError(t, err) + svc := NewServiceMemory(log) + ctx := context.Background() + _, err = svc.QueueMessageWithMetadata(ctx, "session-typed", "task", "content", "", QueuedByUser, false, nil, + map[string]interface{}{ + "tags": []string{"original"}, + "labels": map[string]string{"name": "original"}, + }, + ) + require.NoError(t, err) + + status := svc.GetStatus(ctx, "session-typed") + require.Len(t, status.Entries, 1) + status.Entries[0].Metadata["tags"].([]string)[0] = "mutated" + status.Entries[0].Metadata["labels"].(map[string]string)["name"] = "mutated" + + fresh := svc.GetStatus(ctx, "session-typed") + require.Equal(t, []string{"original"}, fresh.Entries[0].Metadata["tags"]) + require.Equal(t, map[string]string{"name": "original"}, fresh.Entries[0].Metadata["labels"]) +} +func TestSnapshotSessionReturnsDetachedEntrySnapshots(t *testing.T) { + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console", OutputPath: "stderr"}) + require.NoError(t, err) + svc := NewServiceMemory(log) + ctx := context.Background() + _, err = svc.QueueMessageWithMetadata(ctx, "session-snapshot", "task", "content", "", QueuedByUser, false, + []MessageAttachment{{AttachmentID: "attachment-1", Name: "original.txt"}}, + map[string]interface{}{"nested": map[string]interface{}{"value": "original"}}, + ) + require.NoError(t, err) + + entries, _, err := svc.SnapshotSession(ctx, "session-snapshot") + require.NoError(t, err) + require.Len(t, entries, 1) + entries[0].Attachments[0].Name = "mutated.txt" + entries[0].Metadata["nested"].(map[string]interface{})["value"] = "mutated" + + fresh, _, err := svc.SnapshotSession(ctx, "session-snapshot") + require.NoError(t, err) + require.Equal(t, "original.txt", fresh[0].Attachments[0].Name) + require.Equal(t, "original", fresh[0].Metadata["nested"].(map[string]interface{})["value"]) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/service_test.go b/apps/backend/internal/orchestrator/messagequeue/service_test.go index e01b4885b8..a8c08bd876 100644 --- a/apps/backend/internal/orchestrator/messagequeue/service_test.go +++ b/apps/backend/internal/orchestrator/messagequeue/service_test.go @@ -675,6 +675,18 @@ func TestRemoveEntry(t *testing.T) { assert.ErrorIs(t, err, ErrEntryNotFound) }) + t.Run("invalidates an edit lease for the removed entry", func(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + entry, err := svc.QueueMessage(ctx, "s", "t", "editable", "", QueuedByUser, false, nil) + require.NoError(t, err) + lease, err := svc.BeginEdit(ctx, "s", entry.ID, "connection-a") + require.NoError(t, err) + + require.NoError(t, svc.RemoveEntry(ctx, "s", entry.ID)) + assert.ErrorIs(t, svc.EndEdit(ctx, "s", entry.ID, lease.LeaseID, "connection-a"), ErrEditLeaseNotFound) + }) + t.Run("rejects deletion from a foreign session", func(t *testing.T) { svc := setupService(t) ctx := context.Background() @@ -720,7 +732,7 @@ func TestRemoveEntry(t *testing.T) { reserved, ok := svc.ReserveQueued(ctx, "s") require.True(t, ok) assert.ErrorIs(t, svc.RemoveEntry(ctx, "s", reserved.ID), ErrEntryNotFound) - require.NoError(t, svc.AcknowledgeQueued(ctx, "s", reserved.ID)) + require.NoError(t, svc.AcknowledgeQueued(ctx, reserved)) }) } @@ -762,7 +774,7 @@ func TestCancelAllPreservesDurableEntryReservedInFlight(t *testing.T) { require.NoError(t, err) assert.Equal(t, 3, removed) assert.Equal(t, 0, svc.GetStatus(ctx, "s").Count) - require.NoError(t, svc.AcknowledgeQueued(ctx, "s", reserved.ID)) + require.NoError(t, svc.AcknowledgeQueued(ctx, reserved)) } func TestGetStatus(t *testing.T) { @@ -815,7 +827,10 @@ func TestGetStatus(t *testing.T) { require.NoError(t, svc.RequeueAtHead(ctx, reserved)) assert.Equal(t, 1, svc.GetStatus(ctx, "s").Count) - require.NoError(t, svc.AcknowledgeQueued(ctx, "s", reserved.ID)) + require.ErrorIs(t, svc.AcknowledgeQueued(ctx, reserved), ErrLifecycleReservationChanged) + current, ok := svc.ReserveQueued(ctx, "s") + require.True(t, ok) + require.NoError(t, svc.AcknowledgeQueued(ctx, current)) assert.Equal(t, 0, svc.GetStatus(ctx, "s").Count) }) @@ -836,6 +851,38 @@ func TestGetStatus(t *testing.T) { assert.False(t, taken.IsReservedLifecycleDelivery()) }) } +func TestTakeQueuedDoesNotDeleteReservedLifecycleHead(t *testing.T) { + svc := setupService(t) + ctx := context.Background() + + _, _, accepted, err := svc.QueueLifecycleMessageWithCoalesceKey( + ctx, "s", "t", "pr merged", "", QueuedByWorkflow, false, nil, + map[string]interface{}{"origin": "github_pr_automation"}, + "github-pr:repo:1:merged", true, + ) + require.NoError(t, err) + require.True(t, accepted) + + reserved, ok := svc.ReserveQueued(ctx, "s") + require.True(t, ok) + require.NotNil(t, reserved) + + taken, ok := svc.TakeQueued(ctx, "s") + assert.False(t, ok) + assert.Nil(t, taken) + + restored, ok, err := svc.TakeQueuedEntry(ctx, "s", reserved.ID) + require.NoError(t, err) + require.True(t, ok, "destructive head cleanup must not consume a reserved lifecycle row") + require.Equal(t, reserved.ID, restored.ID) +} +func TestQueuedMessageLegacyGitLabLifecycleOriginIsDurable(t *testing.T) { + message := &QueuedMessage{Metadata: map[string]interface{}{ + "origin": "gitlab_mr_automation", + }} + + require.True(t, message.IsDurableLifecycle()) +} func TestTransferSession(t *testing.T) { t.Run("moves entries and pending move", func(t *testing.T) { diff --git a/apps/backend/internal/orchestrator/messagequeue/service_transfer_compensation_test.go b/apps/backend/internal/orchestrator/messagequeue/service_transfer_compensation_test.go new file mode 100644 index 0000000000..e6ba207c72 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/service_transfer_compensation_test.go @@ -0,0 +1,760 @@ +package messagequeue + +import ( + "context" + "errors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "testing" +) + +func TestTransferSessionWithPreparationRollsBackPreparationError(t *testing.T) { + ctx := context.Background() + service := setupService(t) + queued, err := service.QueueMessage(ctx, "session-old", "task", "handoff", "", QueuedByUser, false, nil) + require.NoError(t, err) + prepareErr := errors.New("preparation failed after side effect") + rollbackCalls := 0 + + err = service.TransferSessionWithPreparation( + ctx, + "session-old", + "session-new", + func(context.Context) error { + return prepareErr + }, + func(context.Context) error { + rollbackCalls++ + return nil + }, + ) + + require.ErrorIs(t, err, prepareErr) + assert.Equal(t, 1, rollbackCalls) + status := service.GetStatus(ctx, "session-old") + require.Len(t, status.Entries, 1) + assert.Equal(t, queued.ID, status.Entries[0].ID) + assert.Empty(t, service.GetStatus(ctx, "session-new").Entries) +} + +func TestDurableSessionTransferPersistsBeforeExternalPreparation(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + service := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + _, err := service.QueueMessage( + ctx, + "session-old", + "task", + "handoff", + "", + QueuedByUser, + false, + []MessageAttachment{{AttachmentID: "attachment"}}, + ) + require.NoError(t, err) + prepareErr := errors.New("external transfer failed") + rollbackErr := errors.New("external rollback failed") + + err = service.TransferSessionWithDurablePreparation( + ctx, + "task", + "session-old", + "session-new", + func(context.Context) error { + compensations, listErr := service.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + require.Len(t, compensations, 1) + assert.Equal(t, "session-old", compensations[0].FromSessionID) + return prepareErr + }, + func(context.Context) error { + return rollbackErr + }, + ) + + require.ErrorIs(t, err, prepareErr) + require.ErrorIs(t, err, rollbackErr) + compensations, listErr := service.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + require.Len(t, compensations, 1) + assert.Equal(t, "session-new", compensations[0].ToSessionID) +} + +func TestDurableSessionTransferIncludesInFlightOrdinaryAttachment(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + service := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + queued, err := service.QueueMessage( + ctx, + "session-old", + "task", + "handoff", + "", + QueuedByUser, + false, + []MessageAttachment{{AttachmentID: "attachment"}}, + ) + require.NoError(t, err) + reserved, ok := service.ReserveQueued(ctx, queued.SessionID) + require.True(t, ok) + require.Equal(t, queued.ID, reserved.ID) + assert.Empty(t, service.GetStatus(ctx, queued.SessionID).Entries) + prepareCalls := 0 + + err = service.TransferSessionWithDurablePreparation( + ctx, + queued.TaskID, + queued.SessionID, + "session-new", + func(context.Context) error { + prepareCalls++ + compensations, listErr := service.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + require.Len(t, compensations, 1) + assert.Equal(t, []string{queued.ID}, compensations[0].EntryIDs) + return nil + }, + nil, + ) + + require.NoError(t, err) + assert.Equal(t, 1, prepareCalls) + pending, err := service.ListPendingQueueDispatches(ctx) + require.NoError(t, err) + require.Len(t, pending, 1) + assert.Equal(t, "session-new", pending[0].Message.SessionID) + assert.Equal(t, queued.ID, pending[0].Message.ID) +} + +func TestDurableSessionTransferIncludesCleanupOnlyAttachmentClaim(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + service := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + cleanup := AttachmentCleanup{ + SessionID: "session-cleanup-only-old", EntryID: "entry-cleanup-only", + OperationID: "operation-cleanup-only", TaskID: "task-cleanup-only", + Attachments: []MessageAttachment{{AttachmentID: "attachment-cleanup-only"}}, + } + require.NoError(t, service.UpsertAttachmentCleanup(ctx, cleanup)) + prepareCalls := 0 + + err := service.TransferSessionWithDurablePreparation( + ctx, + cleanup.TaskID, + cleanup.SessionID, + "session-cleanup-only-new", + func(context.Context) error { + prepareCalls++ + compensations, listErr := service.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + require.Len(t, compensations, 1) + assert.Equal(t, []string{cleanup.EntryID}, compensations[0].EntryIDs) + return nil + }, + nil, + ) + + require.NoError(t, err) + assert.Equal(t, 1, prepareCalls) + stored, err := service.GetAttachmentCleanup( + ctx, cleanup.SessionID, cleanup.EntryID, cleanup.OperationID, + ) + require.NoError(t, err) + require.NotNil(t, stored) + assert.Equal(t, "session-cleanup-only-new", stored.CurrentSessionID) +} + +func TestDurableSessionTransferFencesConcurrentSourceInsert(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + service := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + _, err := service.QueueMessage( + ctx, + "session-old", + "task", + "handoff", + "", + QueuedByUser, + false, + []MessageAttachment{{AttachmentID: "attachment"}}, + ) + require.NoError(t, err) + prepared := make(chan struct{}) + release := make(chan struct{}) + transferDone := make(chan error, 1) + go func() { + transferDone <- service.TransferSessionWithDurableAttachmentPreparation( + ctx, + "task", + "session-old", + "session-new", + func(context.Context, []string) error { + close(prepared) + <-release + return nil + }, + nil, + ) + }() + <-prepared + + err = repo.Insert(ctx, &QueuedMessage{ + ID: "concurrent", SessionID: "session-old", TaskID: "task", + Content: "concurrent", QueuedBy: QueuedByUser, + }, DefaultMaxPerSession) + require.ErrorContains(t, err, "session transfer in progress") + close(release) + require.NoError(t, <-transferDone) +} + +func TestAttachmentCleanupFollowsActiveSessionTransfer(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + persistence := repo.(interface { + UpsertSessionTransferCompensation(context.Context, SessionTransferCompensation) error + UpsertAttachmentCleanup(context.Context, AttachmentCleanup) error + GetAttachmentCleanup(context.Context, string, string, string) (*AttachmentCleanup, error) + }) + require.NoError(t, persistence.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: "task", FromSessionID: "session-old", ToSessionID: "session-new", + })) + cleanup := AttachmentCleanup{ + SessionID: "session-old", CurrentSessionID: "session-old", + EntryID: "entry", OperationID: "operation", TaskID: "task", + Attachments: []MessageAttachment{{AttachmentID: "attachment"}}, + } + + require.NoError(t, persistence.UpsertAttachmentCleanup(ctx, cleanup)) + stored, err := persistence.GetAttachmentCleanup( + ctx, cleanup.SessionID, cleanup.EntryID, cleanup.OperationID, + ) + require.NoError(t, err) + require.NotNil(t, stored) + assert.Equal(t, "session-new", stored.CurrentSessionID) +} + +func TestAttachmentCleanupDeleteWaitsForActiveSessionTransfer(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + persistence := repo.(interface { + UpsertSessionTransferCompensation(context.Context, SessionTransferCompensation) error + UpsertAttachmentCleanup(context.Context, AttachmentCleanup) error + DeleteAttachmentCleanup(context.Context, string, string, string) error + GetAttachmentCleanup(context.Context, string, string, string) (*AttachmentCleanup, error) + }) + cleanup := AttachmentCleanup{ + SessionID: "session-old", EntryID: "entry", OperationID: "operation", TaskID: "task", + Attachments: []MessageAttachment{{AttachmentID: "attachment"}}, + } + require.NoError(t, persistence.UpsertAttachmentCleanup(ctx, cleanup)) + require.NoError(t, persistence.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: "task", FromSessionID: "session-old", ToSessionID: "session-new", + })) + + err := persistence.DeleteAttachmentCleanup( + ctx, cleanup.SessionID, cleanup.EntryID, cleanup.OperationID, + ) + require.ErrorIs(t, err, ErrSessionTransferInProgress) + stored, err := persistence.GetAttachmentCleanup( + ctx, cleanup.SessionID, cleanup.EntryID, cleanup.OperationID, + ) + require.NoError(t, err) + require.NotNil(t, stored) +} + +func TestActiveSessionTransferFencesExistingQueueMutations(t *testing.T) { + tests := []struct { + name string + mutate func(context.Context, Repository, *QueuedMessage, *QueuedMessage) error + }{ + { + name: "edit", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + return repo.UpdateContent(ctx, first.SessionID, first.ID, "edited", nil, QueuedByUser) + }, + }, + { + name: "purge", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + _, err := repo.PurgeSession(ctx, first.SessionID) + return err + }, + }, + { + name: "purge task", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + _, err := repo.PurgeTask(ctx, first.TaskID) + return err + }, + }, + { + name: "merge", + mutate: func(ctx context.Context, repo Repository, _, second *QueuedMessage) error { + _, err := repo.MergeIntoAbove(ctx, second.SessionID, second.ID, QueuedByUser) + return err + }, + }, + { + name: "reorder", + mutate: func(ctx context.Context, repo Repository, first, second *QueuedMessage) error { + return repo.ReorderEntries(ctx, first.SessionID, []string{second.ID, first.ID}) + }, + }, + { + name: "take pending move", + mutate: func(ctx context.Context, repo Repository, first, _ *QueuedMessage) error { + _, err := repo.TakePendingMove(ctx, first.SessionID) + return err + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + first := insertTestEntry(t, repo, "session-old", "task", "first", QueuedByUser, nil, nil) + second := insertTestEntry(t, repo, "session-old", "task", "second", QueuedByUser, nil, nil) + require.NoError(t, repo.SetPendingMove(ctx, first.SessionID, &PendingMove{ + MoveID: "move", TaskID: "task", WorkflowID: "workflow", WorkflowStepID: "step", + })) + persistence := repo.(interface { + UpsertSessionTransferCompensation(context.Context, SessionTransferCompensation) error + }) + require.NoError(t, persistence.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: "task", FromSessionID: first.SessionID, ToSessionID: "session-new", + })) + + err := test.mutate(ctx, repo, first, second) + + require.ErrorIs(t, err, ErrSessionTransferInProgress) + }) + } +} + +func TestActiveSessionTransferRejectsNonOwnerQueueTransfer(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + entry := insertTestEntry(t, repo, "session-old", "task", "first", QueuedByUser, nil, nil) + persistent := repo.(*sqliteRepository) + require.NoError(t, persistent.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-owner", + TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + })) + + err := repo.TransferSession(ctx, entry.SessionID, "session-third") + + require.ErrorIs(t, err, ErrSessionTransferInProgress) + stored, listErr := repo.ListBySession(ctx, entry.SessionID) + require.NoError(t, listErr) + require.Len(t, stored, 1) + assert.Equal(t, entry.ID, stored[0].ID) +} + +func TestActiveSessionTransferFencesDispatchSettlement(t *testing.T) { + tests := []struct { + name string + settle func(context.Context, *sqliteRepository, *QueuedMessage) error + }{ + {name: "mark accepted", settle: func(ctx context.Context, repo *sqliteRepository, msg *QueuedMessage) error { + return repo.MarkPendingQueueDispatchAccepted(ctx, msg) + }}, + {name: "delete", settle: func(ctx context.Context, repo *sqliteRepository, msg *QueuedMessage) error { + return repo.DeletePendingQueueDispatch(ctx, msg) + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + entry := insertTestEntry(t, repository, "session-old", "task", "first", QueuedByUser, nil, nil) + reserved, err := repository.ReserveHead(ctx, entry.SessionID) + require.NoError(t, err) + require.NotNil(t, reserved) + require.NoError(t, persistent.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-owner", + TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + })) + + err = test.settle(ctx, persistent, reserved) + + require.ErrorIs(t, err, ErrSessionTransferInProgress) + pending, listErr := persistent.ListPendingQueueDispatches(ctx) + require.NoError(t, listErr) + require.Len(t, pending, 1) + assert.False(t, pending[0].Accepted) + }) + } +} + +func TestActiveSessionTransferFencesSendNowSettlement(t *testing.T) { + tests := []struct { + name string + settle func(context.Context, *sqliteRepository, *SendNowClaim) error + }{ + {name: "mark accepted", settle: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.MarkPendingSendNowClaimAccepted(ctx, claim) + }}, + {name: "delete", settle: func(ctx context.Context, repo *sqliteRepository, claim *SendNowClaim) error { + return repo.DeletePendingSendNowClaim(ctx, claim) + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + entry := insertTestEntry(t, repository, "session-old", "task", "first", QueuedByUser, nil, nil) + claim, err := repository.ClaimSendNow(ctx, entry.SessionID, []QueuedMessage{*entry}) + require.NoError(t, err) + require.NoError(t, persistent.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-owner", + TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + })) + + err = test.settle(ctx, persistent, claim) + + require.ErrorIs(t, err, ErrSessionTransferInProgress) + pending, listErr := persistent.ListPendingSendNowClaims(ctx) + require.NoError(t, listErr) + require.Len(t, pending, 1) + assert.False(t, pending[0].Accepted) + }) + } +} + +func setupCrossRepositoryEditTransferFence( + t *testing.T, + withLease bool, +) (*Service, *QueuedMessage, *QueueEditLease) { + t.Helper() + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + service := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := service.QueueMessage( + ctx, "session-old", "task", "first", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + var lease *QueueEditLease + if withLease { + lease, err = service.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + } + require.NoError(t, persistent.UpsertSessionTransferCompensation(ctx, SessionTransferCompensation{ + OperationID: "transfer-owner", + TaskID: entry.TaskID, FromSessionID: entry.SessionID, ToSessionID: "session-new", + })) + return service, entry, lease +} + +func TestActiveSessionTransferFencesEditLeaseLifecycle(t *testing.T) { + ctx := context.Background() + t.Run("begin", func(t *testing.T) { + service, entry, _ := setupCrossRepositoryEditTransferFence(t, false) + lease, err := service.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.ErrorIs(t, err, ErrSessionTransferInProgress) + assert.Nil(t, lease) + }) + t.Run("renew", func(t *testing.T) { + service, entry, lease := setupCrossRepositoryEditTransferFence(t, true) + renewed, err := service.RenewEdit( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection", + ) + require.ErrorIs(t, err, ErrSessionTransferInProgress) + assert.Nil(t, renewed) + stored, getErr := service.GetEditLease(ctx, entry.SessionID, entry.ID) + require.NoError(t, getErr) + assert.Equal(t, lease.LeaseGeneration, stored.LeaseGeneration) + }) + t.Run("end", func(t *testing.T) { + service, entry, lease := setupCrossRepositoryEditTransferFence(t, true) + err := service.EndEdit(ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection") + require.ErrorIs(t, err, ErrSessionTransferInProgress) + _, getErr := service.GetEditLease(ctx, entry.SessionID, entry.ID) + require.NoError(t, getErr) + }) +} + +func TestRemoteSessionTransferRejectsStaleEditLeaseRenewal(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + transferService := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + editService := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-old", "task", "first", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + lease, err := editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + require.NoError(t, transferService.TransferSessionWithDurablePreparation( + ctx, entry.TaskID, entry.SessionID, "session-new", nil, nil, + )) + + renewed, err := editService.RenewEdit( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "connection", + ) + + require.ErrorIs(t, err, ErrEditLeaseNotFound) + assert.Nil(t, renewed) +} + +func TestRemoteSessionTransferRejectsDuplicateEditReplay(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + transferService := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + editService := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-old", "task", "before", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + lease, err := editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + finalizeErr := errors.New("finalize failed") + finalizeCalls := 0 + finalize := func(context.Context, *QueuedMessage) error { + finalizeCalls++ + if finalizeCalls == 1 { + return finalizeErr + } + return nil + } + _, err = editService.UpdateMessageWithLeaseAfterValidationAndFinalize( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation", "connection", + lease.TargetRevision, "after", nil, nil, nil, nil, finalize, + ) + require.ErrorIs(t, err, finalizeErr) + require.NoError(t, transferService.TransferSessionWithDurablePreparation( + ctx, entry.TaskID, entry.SessionID, "session-new", nil, nil, + )) + + _, err = editService.UpdateMessageWithLeaseAfterValidationAndFinalize( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation", "connection", + lease.TargetRevision, "after", nil, nil, nil, nil, finalize, + ) + + require.ErrorIs(t, err, ErrEntryNotFound) + assert.Equal(t, 1, finalizeCalls) +} + +func TestRemoteEditLeaseBlocksHeadReservation(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + editService := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + drainService := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-1", "task", "queued", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + _, err = editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + reserved, ok := drainService.ReserveQueued(ctx, entry.SessionID) + + assert.False(t, ok) + assert.Nil(t, reserved) + status := editService.GetStatus(ctx, entry.SessionID) + require.Len(t, status.Entries, 1) + assert.Equal(t, entry.ID, status.Entries[0].ID) +} + +func TestRemoteEditLeaseBlocksTargetedDrain(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + editService := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + drainService := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-1", "task", "queued", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + _, err = editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + taken, ok, err := drainService.TakeQueuedEntry(ctx, entry.SessionID, entry.ID) + + require.ErrorIs(t, err, ErrEditConflict) + assert.False(t, ok) + assert.Nil(t, taken) +} + +func TestRemoteEditLeaseBlocksSendNowClaim(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + editService := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + sendNowService := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-1", "task", "queued", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + _, err = editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + claim, err := sendNowService.ClaimSendNow(ctx, entry.SessionID, []QueuedMessage{*entry}) + + require.ErrorIs(t, err, ErrEditConflict) + assert.Nil(t, claim) +} + +func TestRemoteEditLeaseBlocksEntryRemoval(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + editService := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + removeService := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-1", "task", "queued", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + _, err = editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + removed, err := removeService.RemoveEntryWithEntry(ctx, entry.SessionID, entry.ID) + + require.ErrorIs(t, err, ErrEditConflict) + assert.Nil(t, removed) +} + +func TestRemoteDisconnectReleasesDurableEditLease(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t) + persistent := repository.(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(persistent.db, persistent.ro) + require.NoError(t, err) + firstService := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + secondService := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := firstService.QueueMessage( + ctx, "session-1", "task", "queued", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + _, err = firstService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + require.Equal(t, 1, firstService.ReleaseEditLeasesForConnection("connection")) + _, err = secondService.BeginEdit(ctx, entry.SessionID, entry.ID, "new-connection") + + require.NoError(t, err) +} + +func TestRemoteEditLeaseBlocksUncoveredMutations(t *testing.T) { + t.Run("update", func(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t).(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(repository.db, repository.ro) + require.NoError(t, err) + owner := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + other := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := owner.QueueMessage(ctx, "session-1", "task", "before", "", QueuedByUser, false, nil) + require.NoError(t, err) + _, err = owner.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + err = other.UpdateMessageWithMetadata(ctx, entry.SessionID, entry.ID, "after", nil, nil, QueuedByUser) + + require.ErrorIs(t, err, ErrEditConflict) + }) + + t.Run("append", func(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t).(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(repository.db, repository.ro) + require.NoError(t, err) + owner := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + other := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, err := owner.QueueMessage(ctx, "session-1", "task", "before", "", QueuedByUser, false, nil) + require.NoError(t, err) + _, err = owner.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + _, _, err = other.AppendContent(ctx, entry.SessionID, entry.TaskID, "after", "", QueuedByUser, false, nil) + + require.ErrorIs(t, err, ErrEditConflict) + }) + + t.Run("coalesce replacement", func(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t).(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(repository.db, repository.ro) + require.NoError(t, err) + owner := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + other := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + entry, _, err := owner.QueueMessageWithCoalesceKey(ctx, "session-1", "task", "before", "", QueuedByUser, false, nil, nil, "key", true) + require.NoError(t, err) + lease, err := owner.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + + replacement, replaced, err := other.QueueMessageWithCoalesceKey(ctx, entry.SessionID, entry.TaskID, "after", "", QueuedByUser, false, nil, nil, "key", true) + require.NoError(t, err) + require.True(t, replaced) + require.Equal(t, entry.ID, replacement.ID) + _, err = owner.UpdateMessageWithLease(ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation", "connection", lease.TargetRevision, "stale", nil, nil) + require.ErrorIs(t, err, ErrEditLeaseNotFound) + }) + + t.Run("automatic merge", func(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t).(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(repository.db, repository.ro) + require.NoError(t, err) + owner := newAutoMergeTestServiceWithRepository(t, repository, DefaultMaxPerSession) + other := newAutoMergeTestServiceWithRepository(t, secondRepository, DefaultMaxPerSession) + owner.SetAutoMergeEnabled(false) + first, err := owner.QueueMessage(ctx, "session-1", "task", "first", "", QueuedByUser, false, nil) + require.NoError(t, err) + second, err := owner.QueueMessage(ctx, first.SessionID, first.TaskID, "second", "", QueuedByUser, false, nil) + require.NoError(t, err) + _, err = owner.BeginEdit(ctx, second.SessionID, second.ID, "connection") + require.NoError(t, err) + other.SetAutoMergeEnabled(true) + + queued, err := other.QueueMessage(ctx, first.SessionID, first.TaskID, "third", "", QueuedByUser, false, nil) + + require.NoError(t, err) + require.NotEqual(t, first.ID, queued.ID) + status := owner.GetStatus(ctx, first.SessionID) + require.Len(t, status.Entries, 3) + require.Equal(t, "first", status.Entries[0].Content) + }) + + t.Run("full queue candidate merge", func(t *testing.T) { + ctx := context.Background() + repository := newTestSQLiteRepo(t).(*sqliteRepository) + secondRepository, err := NewSQLiteRepository(repository.db, repository.ro) + require.NoError(t, err) + owner := newAutoMergeTestServiceWithRepository(t, repository, 1) + other := newAutoMergeTestServiceWithRepository(t, secondRepository, 1) + first, err := owner.QueueMessage(ctx, "session-1", "task", "first", "", QueuedByUser, false, nil) + require.NoError(t, err) + _, err = owner.BeginEdit(ctx, first.SessionID, first.ID, "connection") + require.NoError(t, err) + + _, err = other.QueueMessage(ctx, first.SessionID, first.TaskID, "second", "", QueuedByUser, false, nil) + + require.ErrorIs(t, err, ErrQueueFull) + status := owner.GetStatus(ctx, first.SessionID) + require.Len(t, status.Entries, 1) + require.Equal(t, "first", status.Entries[0].Content) + }) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/service_transfer_lease_test.go b/apps/backend/internal/orchestrator/messagequeue/service_transfer_lease_test.go new file mode 100644 index 0000000000..1d146312b4 --- /dev/null +++ b/apps/backend/internal/orchestrator/messagequeue/service_transfer_lease_test.go @@ -0,0 +1,309 @@ +package messagequeue + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSessionTransferCompensationRejectsActiveReplacement(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + persistence := repo.(interface { + UpsertSessionTransferCompensation(context.Context, SessionTransferCompensation) error + DeleteSessionTransferCompensation(context.Context, string, string, string, string) error + ListSessionTransferCompensations(context.Context) ([]SessionTransferCompensation, error) + }) + first := SessionTransferCompensation{ + OperationID: "transfer-first", + TaskID: "task", FromSessionID: "session-old", ToSessionID: "session-new", + EntryIDs: []string{"entry-first"}, + } + require.NoError(t, persistence.UpsertSessionTransferCompensation(ctx, first)) + + replacement := first + replacement.OperationID = "transfer-replacement" + replacement.EntryIDs = []string{"entry-replacement"} + err := persistence.UpsertSessionTransferCompensation(ctx, replacement) + + require.ErrorIs(t, err, ErrSessionTransferInProgress) + stored, listErr := persistence.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + require.Len(t, stored, 1) + assert.Equal(t, first.EntryIDs, stored[0].EntryIDs) + first.EntryIDs = []string{"entry-final"} + require.NoError(t, persistence.UpsertSessionTransferCompensation(ctx, first)) + stored, listErr = persistence.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + require.Len(t, stored, 1) + assert.Equal(t, first.EntryIDs, stored[0].EntryIDs) + err = persistence.DeleteSessionTransferCompensation( + ctx, replacement.OperationID, first.TaskID, first.FromSessionID, first.ToSessionID, + ) + require.ErrorIs(t, err, ErrSessionTransferOwnershipLost) + stored, listErr = persistence.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + require.Len(t, stored, 1) + assert.Equal(t, first.OperationID, stored[0].OperationID) + require.NoError(t, persistence.DeleteSessionTransferCompensation( + ctx, first.OperationID, first.TaskID, first.FromSessionID, first.ToSessionID, + )) + stored, listErr = persistence.ListSessionTransferCompensations(ctx) + require.NoError(t, listErr) + assert.Empty(t, stored) +} + +func TestMaintainSessionTransferCompensationLeaseStopIsIdempotent(t *testing.T) { + ctx := context.Background() + service := setupService(t) + compensation := SessionTransferCompensation{ + OperationID: "transfer-operation", + TaskID: "task", + FromSessionID: "session-old", + ToSessionID: "session-new", + } + stopCtx, stop := service.MaintainSessionTransferCompensationLease( + ctx, compensation, compensation.OperationID, + ) + require.NoError(t, stopCtx.Err()) + require.NoError(t, stop()) + require.NoError(t, stop()) +} + +func TestDurableSessionTransferRetainsCompensationAfterLeaseLoss(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + sqlRepo := repo.(*sqliteRepository) + service := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + service.sessionTransferCompensationLeaseRenewInterval = time.Millisecond + _, err := service.QueueMessage( + ctx, + "session-old", + "task", + "handoff", + "", + QueuedByUser, + false, + []MessageAttachment{{AttachmentID: "attachment"}}, + ) + require.NoError(t, err) + preparationStarted := make(chan struct{}) + transferDone := make(chan error, 1) + go func() { + transferDone <- service.TransferSessionWithDurableAttachmentPreparation( + ctx, + "task", + "session-old", + "session-new", + func(prepareCtx context.Context, _ []string) error { + close(preparationStarted) + <-prepareCtx.Done() + return prepareCtx.Err() + }, + func(context.Context, []string) error { + return errors.New("rollback blocked") + }, + ) + }() + <-preparationStarted + compensations, err := service.ListSessionTransferCompensations(ctx) + require.NoError(t, err) + require.Len(t, compensations, 1) + _, err = sqlRepo.db.Exec(` + UPDATE queue_session_transfer_compensations + SET recovery_owner = 'another-owner' + WHERE operation_id = ? + `, compensations[0].OperationID) + require.NoError(t, err) + + select { + case err = <-transferDone: + require.Error(t, err) + assert.ErrorContains(t, err, "session transfer lease lost") + case <-time.After(2 * time.Second): + t.Fatal("transfer did not stop after lease loss") + } + compensations, err = service.ListSessionTransferCompensations(ctx) + require.NoError(t, err) + require.Len(t, compensations, 1) +} +func TestDurableSessionTransferKeepsCompensationAfterLeaseLossAndRollback(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + sqlRepo := repo.(*sqliteRepository) + service := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + service.sessionTransferCompensationLeaseRenewInterval = time.Millisecond + entry, err := service.QueueMessage( + ctx, "session-old", "task", "handoff", "", QueuedByUser, false, + []MessageAttachment{{AttachmentID: "attachment"}}, + ) + require.NoError(t, err) + preparationStarted := make(chan struct{}) + transferDone := make(chan error, 1) + go func() { + transferDone <- service.TransferSessionWithDurableAttachmentPreparation( + ctx, entry.TaskID, entry.SessionID, "session-new", + func(prepareCtx context.Context, _ []string) error { + close(preparationStarted) + <-prepareCtx.Done() + return prepareCtx.Err() + }, + nil, + ) + }() + <-preparationStarted + compensations, err := service.ListSessionTransferCompensations(ctx) + require.NoError(t, err) + require.Len(t, compensations, 1) + _, err = sqlRepo.db.Exec(` + UPDATE queue_session_transfer_compensations + SET recovery_owner = 'another-owner' + WHERE operation_id = ? + `, compensations[0].OperationID) + require.NoError(t, err) + select { + case err = <-transferDone: + require.Error(t, err) + assert.ErrorContains(t, err, "session transfer lease lost") + case <-time.After(2 * time.Second): + t.Fatal("transfer did not stop after lease loss") + } + compensations, err = service.ListSessionTransferCompensations(ctx) + require.NoError(t, err) + require.Len(t, compensations, 1) +} + +func TestEditLeaseCleanupIsDurableAcrossSessionReplacement(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + sqlRepo := repo.(*sqliteRepository) + otherRepo, err := NewSQLiteRepository(sqlRepo.db, sqlRepo.ro) + require.NoError(t, err) + editService := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + replaceService := newAutoMergeTestServiceWithRepository(t, otherRepo, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-replace", "task-replace", "before", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + lease, err := editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + replacement := *entry + replacement.Content = "after" + require.NoError(t, replaceService.RestoreSession( + ctx, entry.SessionID, []QueuedMessage{replacement}, nil, + )) + _, err = editService.UpdateMessageWithLease( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation", + "connection", lease.TargetRevision, "stale", nil, nil, + ) + require.ErrorIs(t, err, ErrEditLeaseNotFound) +} + +func TestEditLeaseCleanupIsDurableAcrossTaskPurge(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t) + sqlRepo := repo.(*sqliteRepository) + otherRepo, err := NewSQLiteRepository(sqlRepo.db, sqlRepo.ro) + require.NoError(t, err) + editService := newAutoMergeTestServiceWithRepository(t, repo, DefaultMaxPerSession) + purgeService := newAutoMergeTestServiceWithRepository(t, otherRepo, DefaultMaxPerSession) + entry, err := editService.QueueMessage( + ctx, "session-purge", "task-purge", "before", "", QueuedByUser, false, nil, + ) + require.NoError(t, err) + lease, err := editService.BeginEdit(ctx, entry.SessionID, entry.ID, "connection") + require.NoError(t, err) + _, err = purgeService.PurgeTask(ctx, entry.TaskID) + require.NoError(t, err) + replacement := *entry + replacement.Content = "restored" + require.NoError(t, purgeService.RestoreSession( + ctx, entry.SessionID, []QueuedMessage{replacement}, nil, + )) + _, err = editService.UpdateMessageWithLease( + ctx, entry.SessionID, entry.ID, lease.LeaseID, "operation", + "connection", lease.TargetRevision, "stale", nil, nil, + ) + require.ErrorIs(t, err, ErrEditLeaseNotFound) +} + +func TestExpiredSessionTransferLeaseCannotDeleteCompensation(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + compensation := SessionTransferCompensation{ + OperationID: "expired-operation", + TaskID: "task", + FromSessionID: "session-old", + ToSessionID: "session-new", + } + require.NoError(t, repo.UpsertSessionTransferCompensation(ctx, compensation)) + _, err := repo.db.Exec(` + UPDATE queue_session_transfer_compensations + SET recovery_lease_expires_at = ? + WHERE operation_id = ? + `, time.Now().UTC().Add(-time.Second), compensation.OperationID) + require.NoError(t, err) + require.ErrorIs(t, repo.DeleteSessionTransferCompensation( + ctx, compensation.OperationID, compensation.TaskID, + compensation.FromSessionID, compensation.ToSessionID, + ), ErrSessionTransferOwnershipLost) + stored, err := repo.ListSessionTransferCompensations(ctx) + require.NoError(t, err) + require.Len(t, stored, 1) +} +func TestExpiredSessionTransferLeaseCannotRenew(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + compensation := SessionTransferCompensation{ + OperationID: "expired-renew-operation", + TaskID: "task", + FromSessionID: "session-old", + ToSessionID: "session-new", + } + require.NoError(t, repo.UpsertSessionTransferCompensation(ctx, compensation)) + _, err := repo.db.Exec(` + UPDATE queue_session_transfer_compensations + SET recovery_lease_expires_at = ? + WHERE operation_id = ? + `, time.Now().UTC().Add(-time.Second), compensation.OperationID) + require.NoError(t, err) + + require.ErrorIs(t, repo.renewSessionTransferCompensationLease( + ctx, compensation, compensation.OperationID, + ), ErrSessionTransferOwnershipLost) + ownerID, err := repo.claimSessionTransferCompensationRecovery(ctx, compensation) + require.NoError(t, err) + require.NotEqual(t, compensation.OperationID, ownerID) +} + +func TestSessionTransferCommitRejectsExpiredLease(t *testing.T) { + ctx := context.Background() + repo := newTestSQLiteRepo(t).(*sqliteRepository) + compensation := SessionTransferCompensation{ + OperationID: "expired-operation", + TaskID: "task", + FromSessionID: "session-old", + ToSessionID: "session-new", + } + require.NoError(t, repo.UpsertSessionTransferCompensation(ctx, compensation)) + tx, err := repo.beginAuthorizedSessionTransferTx( + ctx, compensation.FromSessionID, compensation.ToSessionID, compensation.OperationID, + ) + require.NoError(t, err) + _, err = tx.Exec(` + UPDATE queue_session_transfer_compensations + SET recovery_lease_expires_at = ? + WHERE operation_id = ? + `, time.Now().UTC().Add(-time.Second), compensation.OperationID) + require.NoError(t, err) + err = commitAuthorizedSessionTransferTx( + ctx, tx, repo.db, compensation.FromSessionID, + compensation.ToSessionID, compensation.OperationID, + ) + require.ErrorIs(t, err, ErrSessionTransferOwnershipLost) + require.NoError(t, tx.Rollback()) +} diff --git a/apps/backend/internal/orchestrator/messagequeue/types.go b/apps/backend/internal/orchestrator/messagequeue/types.go index b798eb7fd5..3548122e9d 100644 --- a/apps/backend/internal/orchestrator/messagequeue/types.go +++ b/apps/backend/internal/orchestrator/messagequeue/types.go @@ -11,8 +11,8 @@ const DefaultMaxPerSession = 10 // Sender identities written to QueuedMessage.QueuedBy. The handlers default // any empty user-supplied identity to QueuedByUser so the UpdateMessage -// ownership guard always runs against a non-empty value. Agent, workflow, and -// server identities are reserved for backend dispatch paths. +// ownership guard always runs against a non-empty value. Agent, workflow, +// server, and move-task identities are reserved for backend dispatch paths. const ( QueuedByUser = "user" QueuedByAgent = "agent" @@ -31,7 +31,7 @@ const ( // additive prompts from one agent into a single delivery. See ADR 0051. func IsReservedQueuedBy(queuedBy string) bool { switch queuedBy { - case QueuedByAgent, QueuedByWorkflow, QueuedByServer: + case QueuedByAgent, QueuedByWorkflow, QueuedByServer, QueuedByMoveTask: return true default: return false @@ -72,6 +72,8 @@ const MetadataLifecycleGeneration = "lifecycle_queue_generation" // but are hidden from pending queue status until acknowledged or released. const MetadataLifecycleReserved = "lifecycle_reserved_in_flight" +const metadataLifecycleReservationID = "lifecycle_reservation_id" + // MetadataLifecycleReservationIncarnation binds a retained row to the // immutable session incarnation that reserved it. A replacement incarnation // may discard that stale reservation but must never deliver it. The historical @@ -123,12 +125,51 @@ var ( // ErrLifecycleCancelled means an archive/delete purge invalidated a // previously accepted lifecycle entry before it could be retried. ErrLifecycleCancelled = errors.New("lifecycle queue entry cancelled") + // ErrQueueDispatchClaimChanged means the durable ordinary-dispatch claim + // was cleared or transferred before its worker attempted to settle it. + ErrQueueDispatchClaimChanged = errors.New("queue dispatch claim changed") + // ErrLifecycleReservationChanged means a newer lifecycle delivery attempt + // replaced the reservation being acknowledged. + ErrLifecycleReservationChanged = errors.New("lifecycle reservation changed") + // ErrSessionTransferInProgress prevents a queue mutation from entering a + // session while its rows and external attachment claims are being rebound. + ErrSessionTransferInProgress = errors.New("session transfer in progress") + // ErrSessionTransferOwnershipLost prevents an older transfer attempt from + // deleting or replacing a newer durable transfer fence. + ErrSessionTransferOwnershipLost = errors.New("session transfer ownership lost") + // ErrEditConflict means a queue entry is currently held by another editor. + ErrEditConflict = errors.New("queue entry edit conflict") + // ErrEditLeaseNotFound means a lease is missing, expired, or owned by + // another WebSocket connection. + ErrEditLeaseNotFound = errors.New("queue edit lease not found") + // ErrEditRevisionConflict means the target changed since edit.begin. + ErrEditRevisionConflict = errors.New("queue edit target revision conflict") // ErrAutoMergePolicyChanged means admission's immutable policy snapshot no // longer matches the durable session override. The caller must re-resolve // policy before deciding whether to fold. ErrAutoMergePolicyChanged = errors.New("queue Auto-merge policy changed") ) +const QueueEditLeaseTTL = 60 * time.Second + +// QueueEditLease is the server-issued, connection-bound hold for one entry. +type QueueEditLease struct { + SessionID string `json:"session_id"` + EntryID string `json:"entry_id"` + LeaseID string `json:"lease_id"` + TargetRevision int64 `json:"target_revision"` + LeaseGeneration int64 `json:"lease_generation,omitempty"` + ExpiresAt time.Time `json:"expires_at,omitempty"` + // connectionID and operation fields are server-side fencing state. + connectionID string + taskID string + lastOperationID string + lastOperationHash string + lastOperationResult int64 + lastOperationPrevious *QueuedMessage + lastOperationFinalized bool +} + // QueueSessionIdentity is the immutable authority for session-scoped queue work. type QueueSessionIdentity struct { TaskID string `json:"task_id"` @@ -180,7 +221,21 @@ type QueuedMessage struct { // reservedLifecycleDelivery is process-local evidence that ReserveHead // retained this durable row for acknowledgement. reservedLifecycleDelivery bool - reservationIdentity QueueSessionIdentity + + // dispatchAttemptID and lifecycleReservationID identify the exact durable + // delivery attempts this process may settle. They are intentionally absent + // from the public JSON wire shape. + dispatchAttemptID string + lifecycleReservationID string + + // reservationSessionGeneration and reservationLifecycleGeneration fence a + // FIFO dispatch source to the generations observed while it was reserved. + // They are process-local and only used if reservationGenerationsCaptured is + // true. + reservationSessionGeneration int64 + reservationLifecycleGeneration int64 + reservationGenerationsCaptured bool + reservationIdentity QueueSessionIdentity } // QueueRemovalResult is the atomic outcome of a user-driven queue deletion. @@ -191,6 +246,20 @@ type QueueRemovalResult struct { Retained []QueuedMessage } +// ReservationGenerations returns the session and lifecycle generations +// captured when this entry was reserved for FIFO dispatch. +func (m *QueuedMessage) ReservationGenerations() (int64, int64, bool) { + if m == nil { + return 0, 0, false + } + return m.reservationSessionGeneration, m.reservationLifecycleGeneration, m.reservationGenerationsCaptured +} + +const ( + lifecycleOriginGitHubPR = "github_pr_automation" + lifecycleOriginGitLabMR = "gitlab_mr_automation" +) + // IsDurableLifecycle reports whether this entry uses reserve/ack delivery. // The origin fallback keeps lifecycle rows queued by older builds safe across // a rolling restart before the explicit marker was introduced. @@ -202,7 +271,7 @@ func (m *QueuedMessage) IsDurableLifecycle() bool { return true } origin, _ := m.Metadata["origin"].(string) - return origin == "github_pr_automation" || origin == "ci_automation" + return origin == lifecycleOriginGitHubPR || origin == "ci_automation" || origin == lifecycleOriginGitLabMR } // IsReservedInFlight reports whether this row was retained for an in-flight @@ -221,18 +290,19 @@ func (m *QueuedMessage) IsReservedLifecycleDelivery() bool { return m != nil && m.reservedLifecycleDelivery } -// markReservedMetadata returns a copy of metadata carrying the in-flight -// reservation marker. -func markReservedMetadata(metadata map[string]interface{}) map[string]interface{} { - return markReservedMetadataForIncarnation(metadata, "") +// markReservedMetadata returns a copy carrying the exact lifecycle delivery +// attempt that may acknowledge the durable row. +func markReservedMetadata(metadata map[string]interface{}, reservationIDs ...string) map[string]interface{} { + marked := copyMessageMetadata(metadata, 3) + marked[MetadataLifecycleReserved] = true + if len(reservationIDs) > 0 && reservationIDs[0] != "" { + marked[metadataLifecycleReservationID] = reservationIDs[0] + } + return marked } func markReservedMetadataForIncarnation(metadata map[string]interface{}, incarnationID string) map[string]interface{} { - marked := make(map[string]interface{}, len(metadata)+2) - for k, v := range metadata { - marked[k] = v - } - marked[MetadataLifecycleReserved] = true + marked := markReservedMetadata(metadata) if incarnationID != "" { marked[MetadataLifecycleReservationIncarnation] = incarnationID } @@ -247,11 +317,19 @@ func lifecycleReservationIncarnation(metadata map[string]interface{}) string { // clearReservedMetadata removes transient delivery ownership from copies // returned to dispatch or written back for retry. func clearReservedMetadata(metadata map[string]interface{}) map[string]interface{} { - cleared := make(map[string]interface{}, len(metadata)) - for k, v := range metadata { - if k != MetadataLifecycleReserved && k != MetadataLifecycleReservationIncarnation { - cleared[k] = v + cleared := copyMessageMetadata(metadata, 0) + if cleared == nil { + cleared = make(map[string]interface{}) + } + for k := range cleared { + if k != MetadataLifecycleReserved && k != metadataLifecycleReservationID { + if k == MetadataLifecycleReservationIncarnation { + delete(cleared, k) + continue + } + continue } + delete(cleared, k) } return cleared } @@ -267,6 +345,42 @@ type MessageAttachment struct { DeliveryMode string `json:"delivery_mode,omitempty"` } +// AttachmentCleanup records a durable obligation to release attachment claims +// after a queued message no longer references them. +type AttachmentCleanup struct { + SessionID string + CurrentSessionID string + EntryID string + OperationID string + TaskID string + OwnerID string + LeaseID string + RemoveEntry bool + Attachments []MessageAttachment + EntryFingerprint string + ClaimPending bool + CreatedAt time.Time +} + +type AttachmentCleanupLocator struct { + SessionID string `json:"session_id"` + EntryID string `json:"entry_id"` + OperationID string `json:"operation_id"` +} + +// SessionTransferCompensation records an attachment binding that must be +// reconciled with the durable queue location after an interrupted transfer. +type SessionTransferCompensation struct { + OperationID string + TaskID string + FromSessionID string + ToSessionID string + EntryIDs []string + AttachmentIDs []string + CleanupLocators []AttachmentCleanupLocator + CreatedAt time.Time +} + // QueueStatus is the per-session view returned to clients: full ordered list of // pending entries plus capacity info. type QueueStatus struct { diff --git a/apps/backend/internal/orchestrator/queue_dispatch_recovery.go b/apps/backend/internal/orchestrator/queue_dispatch_recovery.go new file mode 100644 index 0000000000..11009a6bf2 --- /dev/null +++ b/apps/backend/internal/orchestrator/queue_dispatch_recovery.go @@ -0,0 +1,48 @@ +package orchestrator + +import ( + "context" + "errors" + "fmt" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "go.uber.org/zap" +) + +func (s *Service) reconcilePendingQueueDispatchesOnStartup(ctx context.Context) error { + if s.messageQueue == nil || !s.messageQueue.PendingQueueDispatchPersistenceAvailable() { + return nil + } + pending, err := s.messageQueue.ListPendingQueueDispatches(ctx) + if err != nil { + return fmt.Errorf("list pending queue dispatches: %w", err) + } + for i := range pending { + dispatch := &pending[i] + msg := &dispatch.Message + if dispatch.Accepted { + if err := s.messageQueue.DeletePendingQueueDispatch(ctx, msg); err != nil { + return fmt.Errorf("acknowledge accepted queue dispatch %s: %w", msg.ID, err) + } + continue + } + _, findErr := s.messageQueue.FindEntryByID(ctx, msg.ID) + switch { + case findErr == nil: + if err := s.messageQueue.DeletePendingQueueDispatch(ctx, msg); err != nil { + return fmt.Errorf("acknowledge restored queue dispatch %s: %w", msg.ID, err) + } + case errors.Is(findErr, messagequeue.ErrEntryNotFound): + s.logger.Warn("restoring queued dispatch with unknown pre-crash delivery outcome", + zap.String("entry_id", msg.ID), + zap.String("session_id", msg.SessionID), + zap.String("task_id", msg.TaskID)) + if _, err := s.messageQueue.RestoreMessage(context.WithoutCancel(ctx), msg); err != nil { + return fmt.Errorf("restore pending queue dispatch %s: %w", msg.ID, err) + } + default: + return fmt.Errorf("locate pending queue dispatch %s: %w", msg.ID, findErr) + } + } + return nil +} diff --git a/apps/backend/internal/orchestrator/queue_dispatch_recovery_test.go b/apps/backend/internal/orchestrator/queue_dispatch_recovery_test.go new file mode 100644 index 0000000000..2590df53a8 --- /dev/null +++ b/apps/backend/internal/orchestrator/queue_dispatch_recovery_test.go @@ -0,0 +1,212 @@ +package orchestrator + +import ( + "context" + "errors" + "testing" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" +) + +func TestOrdinaryQueueDispatchRestoresAfterProcessRestart(t *testing.T) { + ctx := context.Background() + dbPath := t.TempDir() + "/queue.db" + queue, db := newWorkflowTransferQueue(t, dbPath) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := queue.ReserveQueued(ctx, "session-1") + if !ok || reserved.ID != source.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + restarted := &Service{logger: testLogger(), messageQueue: restartedQueue} + reconciler, ok := any(restarted).(interface { + reconcilePendingQueueDispatchesOnStartup(context.Context) error + }) + if !ok { + t.Fatal("orchestrator cannot recover pending ordinary queue dispatches") + } + if err := reconciler.reconcilePendingQueueDispatchesOnStartup(ctx); err != nil { + t.Fatal(err) + } + status := restartedQueue.GetStatus(ctx, "session-1") + if len(status.Entries) != 1 || status.Entries[0].ID != source.ID { + t.Fatalf("queue after ordinary dispatch recovery = %#v, want source %s", status.Entries, source.ID) + } +} + +func TestAcceptedOrdinaryQueueDispatchIsAcknowledgedAfterProcessRestart(t *testing.T) { + ctx := context.Background() + dbPath := t.TempDir() + "/queue.db" + queue, db := newWorkflowTransferQueue(t, dbPath) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := queue.ReserveQueued(ctx, "session-1") + if !ok || reserved.ID != source.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + if err := queue.MarkPendingQueueDispatchAccepted(ctx, reserved); err != nil { + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + restarted := &Service{logger: testLogger(), messageQueue: restartedQueue} + if err := restarted.reconcilePendingQueueDispatchesOnStartup(ctx); err != nil { + t.Fatal(err) + } + status := restartedQueue.GetStatus(ctx, "session-1") + if len(status.Entries) != 0 { + t.Fatalf("queue after accepted ordinary dispatch recovery = %#v, want empty", status.Entries) + } + pending, err := restartedQueue.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("pending dispatches = %#v, want empty", pending) + } +} + +func TestOrdinaryQueueDispatchCallbackPersistsAcceptance(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, t.TempDir()+"/queue.db") + t.Cleanup(func() { _ = db.Close() }) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := queue.ReserveQueued(ctx, "session-1") + if !ok || reserved.ID != source.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + service := &Service{logger: testLogger(), messageQueue: queue} + afterDispatch := service.queuedMessageAfterDispatch(ctx, reserved, false) + if afterDispatch == nil { + t.Fatal("ordinary queue dispatch has no durable acceptance callback") + } + if err := afterDispatch(); err != nil { + t.Fatal(err) + } + pending, err := queue.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || !pending[0].Accepted || pending[0].Message.ID != source.ID { + t.Fatalf("pending dispatches = %#v, want accepted %s", pending, source.ID) + } +} + +func TestClearedOrdinaryQueueDispatchIsNotRestored(t *testing.T) { + ctx := context.Background() + dbPath := t.TempDir() + "/queue.db" + queue, db := newWorkflowTransferQueue(t, dbPath) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := queue.ReserveQueued(ctx, "session-1") + if !ok || reserved.ID != source.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + if _, err := queue.CancelAll(ctx, "session-1"); err != nil { + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + restarted := &Service{logger: testLogger(), messageQueue: restartedQueue} + if err := restarted.reconcilePendingQueueDispatchesOnStartup(ctx); err != nil { + t.Fatal(err) + } + if status := restartedQueue.GetStatus(ctx, "session-1"); len(status.Entries) != 0 { + t.Fatalf("queue after cleared dispatch recovery = %#v, want empty", status.Entries) + } +} + +func TestAcceptedOrdinaryDispatchFailureAcknowledgesRecoveryClaim(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, t.TempDir()+"/queue.db") + t.Cleanup(func() { _ = db.Close() }) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := queue.ReserveQueued(ctx, "session-1") + if !ok || reserved.ID != source.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + service := &Service{logger: testLogger(), messageQueue: queue} + service.finishQueuedMessageExecution( + ctx, "session-1", "session-1", reserved, nil, false, false, + &acceptedPromptDispatchError{err: errors.New("durable acceptance publication failed")}, + ) + pending, err := queue.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("pending dispatches = %#v, want empty", pending) + } +} + +func TestTransferredOrdinaryDispatchRestoresToDestinationAfterRestart(t *testing.T) { + ctx := context.Background() + dbPath := t.TempDir() + "/queue.db" + queue, db := newWorkflowTransferQueue(t, dbPath) + source, err := queue.QueueMessage( + ctx, "session-old", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + reserved, ok := queue.ReserveQueued(ctx, source.SessionID) + if !ok || reserved.ID != source.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + if err := queue.TransferSession(ctx, source.SessionID, "session-new"); err != nil { + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + restarted := &Service{logger: testLogger(), messageQueue: restartedQueue} + if err := restarted.reconcilePendingQueueDispatchesOnStartup(ctx); err != nil { + t.Fatal(err) + } + if status := restartedQueue.GetStatus(ctx, source.SessionID); len(status.Entries) != 0 { + t.Fatalf("source queue after transfer recovery = %#v, want empty", status.Entries) + } + status := restartedQueue.GetStatus(ctx, "session-new") + if len(status.Entries) != 1 || status.Entries[0].ID != source.ID { + t.Fatalf("destination queue after transfer recovery = %#v, want %s", status.Entries, source.ID) + } +} diff --git a/apps/backend/internal/orchestrator/queue_purge_status_test.go b/apps/backend/internal/orchestrator/queue_purge_status_test.go index f13acfcecc..1ecfd6067e 100644 --- a/apps/backend/internal/orchestrator/queue_purge_status_test.go +++ b/apps/backend/internal/orchestrator/queue_purge_status_test.go @@ -92,9 +92,18 @@ func TestDeleteSessionCancelsQueuedPromptsAndPublishesStatus(t *testing.T) { t.Fatalf("Subscribe: %v", err) } + if _, _, accepted, err := svc.messageQueue.QueueLifecycleMessageWithCoalesceKey( + ctx, "session-drop", "task-session-queue", "in-flight lifecycle", "", "workflow", false, nil, + map[string]interface{}{"origin": "github_pr_automation"}, "session-delete-lifecycle", true, + ); err != nil || !accepted { + t.Fatalf("QueueLifecycleMessageWithCoalesceKey: accepted=%v err=%v", accepted, err) + } if _, err := svc.messageQueue.QueueMessage(ctx, "session-drop", "task-session-queue", "orphan me", "", "user", false, nil); err != nil { t.Fatalf("QueueMessage drop: %v", err) } + if _, ok, autoRun := svc.messageQueue.ReserveQueuedWithAutoRun(ctx, "session-drop"); !autoRun || !ok { + t.Fatalf("reserve lifecycle queue entry: auto_run=%v ok=%v", autoRun, ok) + } if _, err := svc.messageQueue.QueueMessage(ctx, "session-keep", "task-session-queue", "keep me", "", "user", false, nil); err != nil { t.Fatalf("QueueMessage keep: %v", err) } @@ -106,6 +115,9 @@ func TestDeleteSessionCancelsQueuedPromptsAndPublishesStatus(t *testing.T) { t.Fatalf("DeleteSession: %v", err) } + if _, ok, autoRun := svc.messageQueue.ReserveQueuedWithAutoRun(ctx, "session-drop"); !autoRun || ok { + t.Fatalf("deleted session retained reserved queue entry: auto_run=%v ok=%v", autoRun, ok) + } if got := svc.messageQueue.GetStatus(ctx, "session-drop").Count; got != 0 { t.Fatalf("session-drop queue count = %d, want 0", got) } diff --git a/apps/backend/internal/orchestrator/queue_send_now.go b/apps/backend/internal/orchestrator/queue_send_now.go index 865dc27a54..c7c8b68df3 100644 --- a/apps/backend/internal/orchestrator/queue_send_now.go +++ b/apps/backend/internal/orchestrator/queue_send_now.go @@ -18,11 +18,13 @@ const ( QueueSendNowScopeEntry = "entry" QueueSendNowScopeAll = "all" - sendNowClaimRecoveryTimeout = 10 * time.Second + sendNowClaimRecoveryTimeout = 10 * time.Second + sendNowWorkerShutdownTimeout = 500 * time.Millisecond ) var ( ErrSendNowConflict = errors.New("send-now operation is already in progress") + ErrSendNowEditConflict = errors.New("send-now entry is being edited") ErrSendNowQueueEmpty = errors.New("send-now queue is empty") ErrSendNowEntryNotFound = errors.New("send-now entry is no longer pending") ErrSendNowQueueChanged = errors.New("send-now queue selection changed") @@ -171,8 +173,22 @@ func (s *Service) sendNowRestoreClaimForReservation( restore := &messagequeue.SendNowClaim{ Identity: reservation.identity, Sources: []messagequeue.QueuedMessage{*reservation.source}, + Dispatch: messagequeue.QueuedMessage{SessionID: reservation.sessionID}, SourceGenerations: make(map[string]int64), } + sessionGeneration, lifecycleGeneration, captured := reservation.source.ReservationGenerations() + if captured { + restore.SessionGeneration = sessionGeneration + if reservation.source.TaskID != "" { + restore.SourceGenerations[reservation.source.TaskID] = lifecycleGeneration + } + return restore, nil + } + sessionGeneration, err := s.messageQueue.SessionGeneration(ctx, reservation.sessionID) + if err != nil { + return nil, ErrSendNowQueueChanged + } + restore.SessionGeneration = sessionGeneration if reservation.source.TaskID == "" { return restore, nil } @@ -350,8 +366,8 @@ func (s *Service) claimAndDispatchSendNow(ctx context.Context, identity *message func mapSendNowClaimError(scope string, err error) error { switch { - case errors.Is(err, messagequeue.ErrSendNowEmpty): - return ErrSendNowQueueChanged + case errors.Is(err, messagequeue.ErrEditConflict): + return ErrSendNowEditConflict case errors.Is(err, messagequeue.ErrSendNowReservationConflict): return ErrSendNowConflict case errors.Is(err, messagequeue.ErrSendNowClaimChanged): @@ -379,40 +395,66 @@ func (s *Service) launchSendNowClaim( if s.sendNowCtx == nil { s.sendNowCtx, s.sendNowCancel = context.WithCancel(context.Background()) } + if s.sendNowWorkers == nil { + s.sendNowWorkers = &sync.WaitGroup{} + } workerCtx := s.sendNowCtx - s.sendNowWorkers.Add(1) + workers := s.sendNowWorkers + workers.Add(1) s.sendNowMu.Unlock() go func() { - defer s.sendNowWorkers.Done() + defer workers.Done() s.executeSendNowClaimWithContext(workerCtx, claim, reservation) }() return true } -func (s *Service) resetSendNowWorkers() { +func (s *Service) resetSendNowWorkers() error { s.sendNowMu.Lock() defer s.sendNowMu.Unlock() + if s.sendNowDrain != nil { + select { + case <-s.sendNowDrain: + default: + return errors.New("previous Send Now workers are still recovering") + } + } s.sendNowStopped = false s.sendNowCtx, s.sendNowCancel = context.WithCancel(context.Background()) + s.sendNowWorkers = &sync.WaitGroup{} + s.sendNowDrain = nil + return nil } func (s *Service) stopSendNowWorkers() { s.sendNowMu.Lock() s.sendNowStopped = true cancel := s.sendNowCancel + workers := s.sendNowWorkers + done := s.sendNowDrain + startWait := workers != nil && done == nil + if startWait { + newDone := make(chan struct{}) + done = newDone + s.sendNowDrain = newDone + } s.sendNowMu.Unlock() if cancel != nil { cancel() } - done := make(chan struct{}) - go func() { - s.sendNowWorkers.Wait() - close(done) - }() + if workers == nil { + return + } + if startWait { + go func() { + workers.Wait() + close(done) + }() + } select { case <-done: - case <-time.After(sendNowClaimRecoveryTimeout): - s.logger.Warn("timed out waiting for send-now workers during shutdown") + case <-time.After(sendNowWorkerShutdownTimeout): + s.logger.Warn("timed out waiting for Send Now workers; recovery remains owned by detached workers") } } @@ -469,6 +511,21 @@ func (s *Service) executeSendNowClaimWithContext( } } if err := s.promptSendNowClaim(ctx, claim); err != nil { + var acceptedDispatch *acceptedPromptDispatchError + if errors.As(err, &acceptedDispatch) { + s.logger.Warn("send-now replacement prompt was accepted but acceptance persistence failed; settling without restore", + zap.String("session_id", sessionID), zap.Error(err)) + if markErr := s.markSendNowClaimAcceptedWithRetry(ctx, claim); markErr != nil { + s.logger.Error("failed to persist accepted send-now queue claim", + zap.String("session_id", sessionID), zap.Error(markErr)) + } + if ackErr := s.acknowledgeSendNowClaimWithRetry(ctx, claim); ackErr != nil { + s.logger.Error("failed to acknowledge accepted send-now queue claim", + zap.String("session_id", sessionID), zap.Error(ackErr)) + } + s.publishQueueStatusEvent(context.Background(), sessionID) + return + } s.logger.Warn("send-now replacement prompt failed; restoring queue claim", zap.String("session_id", sessionID), zap.Error(err)) restore() @@ -515,6 +572,9 @@ func (s *Service) promptSendNowClaim(ctx context.Context, claim *messagequeue.Se _, err := s.promptTask(ctx, claim.Dispatch.TaskID, sessionID, promptContent, claim.Dispatch.Model, claim.Dispatch.PlanMode, attachments, false, promptTaskOptions{ claimEntryID: claim.Dispatch.ID, + afterDispatch: func() error { + return s.markSendNowClaimAcceptedWithRetry(ctx, claim) + }, afterClaim: func() error { if !s.queuedDispatchIdentityIsCurrent(ctx, claim.Identity) { return errLifecyclePromptReservationSuperseded @@ -556,6 +616,17 @@ func (s *Service) acknowledgeSendNowClaimWithRetry( }) } +func (s *Service) markSendNowClaimAcceptedWithRetry( + ctx context.Context, + claim *messagequeue.SendNowClaim, +) error { + if claim == nil || s.messageQueue == nil || !s.messageQueue.PendingSendNowClaimPersistenceAvailable() { + return nil + } + return s.retrySendNowClaimMutation(ctx, func(recoveryCtx context.Context) error { + return s.messageQueue.MarkPendingSendNowClaimAccepted(recoveryCtx, claim) + }) +} func (s *Service) retrySendNowClaimMutation( ctx context.Context, mutate func(context.Context) error, @@ -580,3 +651,35 @@ func (s *Service) retrySendNowClaimMutation( } return err } + +func (s *Service) reconcilePendingSendNowClaimsOnStartup(ctx context.Context) error { + if s.messageQueue == nil || !s.messageQueue.PendingSendNowClaimPersistenceAvailable() { + return nil + } + claims, err := s.messageQueue.ListPendingSendNowClaims(ctx) + if err != nil { + return fmt.Errorf("list pending Send Now claims: %w", err) + } + for i := range claims { + pending := &claims[i] + claim := &pending.Claim + var recoveryErr error + if pending.Accepted { + recoveryErr = s.acknowledgeSendNowClaimWithRetry(ctx, claim) + } else { + recoveryErr = s.restoreSendNowClaimWithRetry(ctx, claim) + } + if recoveryErr != nil { + if !errors.Is(recoveryErr, messagequeue.ErrSendNowClaimChanged) { + return fmt.Errorf("recover pending Send Now claim for session %s: %w", claim.Dispatch.SessionID, recoveryErr) + } + if deleteErr := s.messageQueue.DeletePendingSendNowClaim( + context.WithoutCancel(ctx), claim, + ); deleteErr != nil && !errors.Is(deleteErr, messagequeue.ErrSendNowClaimChanged) { + return fmt.Errorf("discard superseded Send Now claim for session %s: %w", claim.Dispatch.SessionID, deleteErr) + } + } + s.publishQueueStatusEvent(ctx, claim.Dispatch.SessionID) + } + return nil +} diff --git a/apps/backend/internal/orchestrator/queue_send_now_recovery_test.go b/apps/backend/internal/orchestrator/queue_send_now_recovery_test.go new file mode 100644 index 0000000000..114ed86704 --- /dev/null +++ b/apps/backend/internal/orchestrator/queue_send_now_recovery_test.go @@ -0,0 +1,207 @@ +package orchestrator + +import ( + "context" + "database/sql" + "errors" + "path/filepath" + "sync/atomic" + + "github.com/jmoiron/sqlx" + "github.com/kandev/kandev/internal/orchestrator/executor" + "github.com/kandev/kandev/internal/task/models" + "testing" + + "github.com/kandev/kandev/internal/orchestrator/messagequeue" +) + +func TestSendNowOrdinaryClaimRestoresAfterProcessRestart(t *testing.T) { + ctx := context.Background() + dbPath := t.TempDir() + "/queue.db" + queue, db := newWorkflowTransferQueue(t, dbPath) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + if _, err := queue.ClaimSendNow(ctx, "session-1", []messagequeue.QueuedMessage{*source}); err != nil { + t.Fatal(err) + } + if queue.GetStatus(ctx, "session-1").Count != 0 { + t.Fatal("claimed ordinary prompt remained visible before restart") + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + restarted := &Service{logger: testLogger(), messageQueue: restartedQueue} + if err := restarted.reconcilePendingSendNowClaimsOnStartup(ctx); err != nil { + t.Fatal(err) + } + status := restartedQueue.GetStatus(ctx, "session-1") + if len(status.Entries) != 1 || status.Entries[0].ID != source.ID { + t.Fatalf("queue after restart recovery = %#v, want source %s", status.Entries, source.ID) + } +} + +func TestSendNowWorkerCancellationAfterClaimRestoresOrdinarySource(t *testing.T) { + ctx := context.Background() + queue, db := newWorkflowTransferQueue(t, t.TempDir()+"/queue.db") + t.Cleanup(func() { _ = db.Close() }) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "ordinary prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + claim, err := queue.ClaimSendNow(ctx, "session-1", []messagequeue.QueuedMessage{*source}) + if err != nil { + t.Fatal(err) + } + workerCtx, cancel := context.WithCancel(ctx) + cancel() + svc := &Service{logger: testLogger(), messageQueue: queue} + + svc.executeSendNowClaimWithContext(workerCtx, claim, nil) + + status := queue.GetStatus(ctx, "session-1") + if len(status.Entries) != 1 || status.Entries[0].ID != source.ID { + t.Fatalf("queue after cancelled claimed worker = %#v, want source %s", status.Entries, source.ID) + } +} + +func TestSendNowAcceptedClaimIsAcknowledgedAfterProcessRestart(t *testing.T) { + ctx := context.Background() + dbPath := t.TempDir() + "/queue.db" + queue, db := newWorkflowTransferQueue(t, dbPath) + source, err := queue.QueueMessage( + ctx, "session-1", "task-1", "accepted prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + claim, err := queue.ClaimSendNow(ctx, "session-1", []messagequeue.QueuedMessage{*source}) + if err != nil { + t.Fatal(err) + } + marker, ok := any(queue).(interface { + MarkPendingSendNowClaimAccepted(context.Context, *messagequeue.SendNowClaim) error + }) + if !ok { + t.Fatal("message queue cannot durably mark accepted Send Now claims") + } + if err := marker.MarkPendingSendNowClaimAccepted(ctx, claim); err != nil { + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + restartedQueue, restartedDB := newWorkflowTransferQueue(t, dbPath) + t.Cleanup(func() { _ = restartedDB.Close() }) + restarted := &Service{logger: testLogger(), messageQueue: restartedQueue} + if err := restarted.reconcilePendingSendNowClaimsOnStartup(ctx); err != nil { + t.Fatal(err) + } + if status := restartedQueue.GetStatus(ctx, "session-1"); status.Count != 0 { + t.Fatalf("accepted Send Now source was restored after restart: %#v", status.Entries) + } +} + +type transientAcceptedMarkerRepository struct { + messagequeue.Repository + pending interface { + ListPendingSendNowClaims(context.Context) ([]messagequeue.PendingSendNowClaim, error) + MarkPendingSendNowClaimAccepted(context.Context, *messagequeue.SendNowClaim) error + DeletePendingSendNowClaim(context.Context, *messagequeue.SendNowClaim) error + } + failures atomic.Int32 +} + +func (r *transientAcceptedMarkerRepository) ListPendingSendNowClaims( + ctx context.Context, +) ([]messagequeue.PendingSendNowClaim, error) { + return r.pending.ListPendingSendNowClaims(ctx) +} + +func (r *transientAcceptedMarkerRepository) MarkPendingSendNowClaimAccepted( + ctx context.Context, + claim *messagequeue.SendNowClaim, +) error { + if r.failures.Add(-1) >= 0 { + return errors.New("accepted marker temporarily unavailable") + } + return r.pending.MarkPendingSendNowClaimAccepted(ctx, claim) +} + +func (r *transientAcceptedMarkerRepository) DeletePendingSendNowClaim( + ctx context.Context, + claim *messagequeue.SendNowClaim, +) error { + return r.pending.DeletePendingSendNowClaim(ctx, claim) +} + +func TestAcceptedSendNowMarkerFailureDoesNotRestoreDeliveredPrompt(t *testing.T) { + ctx := context.Background() + raw, err := sql.Open("sqlite3", filepath.Join(t.TempDir(), "queue.db")) + if err != nil { + t.Fatal(err) + } + db := sqlx.NewDb(raw, "sqlite3") + t.Cleanup(func() { _ = db.Close() }) + baseRepo, err := messagequeue.NewSQLiteRepository(db, db) + if err != nil { + t.Fatal(err) + } + pending := baseRepo.(interface { + ListPendingSendNowClaims(context.Context) ([]messagequeue.PendingSendNowClaim, error) + MarkPendingSendNowClaimAccepted(context.Context, *messagequeue.SendNowClaim) error + DeletePendingSendNowClaim(context.Context, *messagequeue.SendNowClaim) error + }) + failingRepo := &transientAcceptedMarkerRepository{Repository: baseRepo, pending: pending} + failingRepo.failures.Store(6) + queue := messagequeue.NewService(failingRepo, messagequeue.DefaultMaxPerSession, testLogger()) + + taskRepo := setupTestRepo(t) + seedSession(t, taskRepo, "task-accepted-marker", "session-accepted-marker", "step-1") + seedExecutorRunning(t, taskRepo, "session-accepted-marker", "task-accepted-marker", "exec-1") + session, err := taskRepo.GetTaskSession(ctx, "session-accepted-marker") + if err != nil { + t.Fatal(err) + } + session.State = models.TaskSessionStateWaitingForInput + if err := taskRepo.UpdateTaskSession(ctx, session); err != nil { + t.Fatal(err) + } + agentMgr := &mockAgentManager{isAgentRunning: true, repoForExecutionLookup: taskRepo} + svc := createTestServiceWithAgent(taskRepo, newMockStepGetter(), newMockTaskRepo(), agentMgr) + svc.executor = executor.NewExecutor(agentMgr, taskRepo, testLogger(), executor.ExecutorConfig{}) + svc.messageQueue = queue + + source, err := queue.QueueMessage( + ctx, session.ID, session.TaskID, "accepted prompt", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatal(err) + } + claim, err := queue.ClaimSendNow(ctx, session.ID, []messagequeue.QueuedMessage{*source}) + if err != nil { + t.Fatal(err) + } + svc.markQueuedDispatchInFlight(session.ID, claim.Dispatch.ID) + svc.executeSendNowClaimWithContext(ctx, claim, nil) + + if status := queue.GetStatus(ctx, session.ID); status.Count != 0 { + t.Fatalf("accepted prompt sources were restored: %#v", status.Entries) + } + claims, err := queue.ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatal(err) + } + if len(claims) != 0 { + t.Fatalf("accepted claim was not acknowledged after marker failures: %#v", claims) + } +} diff --git a/apps/backend/internal/orchestrator/queue_send_now_test.go b/apps/backend/internal/orchestrator/queue_send_now_test.go index 6e39d33c9c..7aa8ad706f 100644 --- a/apps/backend/internal/orchestrator/queue_send_now_test.go +++ b/apps/backend/internal/orchestrator/queue_send_now_test.go @@ -72,7 +72,9 @@ func TestSendNowWorkersCanRestartAfterStop(t *testing.T) { t.Fatal("stopping Send Now workers did not mark the worker owner stopped") } - svc.resetSendNowWorkers() + if err := svc.resetSendNowWorkers(); err != nil { + t.Fatal(err) + } if svc.sendNowStopped { t.Fatal("resetting Send Now workers left the worker owner stopped") } @@ -88,6 +90,57 @@ func TestSendNowWorkersCanRestartAfterStop(t *testing.T) { svc.stopSendNowWorkers() } +func TestSendNowRecoveryDetachesWorkerCancellation(t *testing.T) { + svc := &Service{} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + var mutationCtxErr error + + err := svc.retrySendNowClaimMutation(ctx, func(recoveryCtx context.Context) error { + mutationCtxErr = recoveryCtx.Err() + return nil + }) + + if err != nil { + t.Fatalf("detached recovery error = %v", err) + } + if mutationCtxErr != nil { + t.Fatalf("recovery inherited worker cancellation: %v", mutationCtxErr) + } +} + +func TestStopSendNowWorkersReturnsWhenProviderIgnoresCancellation(t *testing.T) { + svc := &Service{logger: testLogger()} + if err := svc.resetSendNowWorkers(); err != nil { + t.Fatal(err) + } + providerStarted := make(chan struct{}) + releaseProvider := make(chan struct{}) + svc.sendNowWorkers.Add(1) + go func() { + defer svc.sendNowWorkers.Done() + close(providerStarted) + <-releaseProvider + }() + <-providerStarted + stopped := make(chan struct{}) + go func() { + svc.stopSendNowWorkers() + close(stopped) + }() + select { + case <-stopped: + case <-time.After(time.Second): + close(releaseProvider) + <-stopped + t.Fatal("Send Now shutdown waited indefinitely for a stuck provider") + } + if err := svc.resetSendNowWorkers(); err == nil { + t.Fatal("restart accepted while the prior Send Now worker still owned recovery") + } + close(releaseProvider) +} + func TestExplicitCancellationDoesNotJoinSendNowOperation(t *testing.T) { operation := &cancelOperation{ done: make(chan struct{}), @@ -331,6 +384,78 @@ func TestSendQueuedNowMissingEntryPreservesAutoRunOff(t *testing.T) { t.Fatalf("rejected Send Now changed queue count to %d", status.Count) } } +func TestSendNowRestoresPendingFIFOWithSessionGeneration(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-1", "session-1", "step-1") + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + + if _, err := svc.messageQueue.QueueMessage( + ctx, "session-1", "task-1", "pending", "", messagequeue.QueuedByUser, false, nil, + ); err != nil { + t.Fatalf("queue message: %v", err) + } + if _, err := svc.messageQueue.PurgeSession(ctx, "session-1"); err != nil { + t.Fatalf("purge session: %v", err) + } + source, err := svc.messageQueue.QueueMessage( + ctx, "session-1", "task-1", "replacement", "", messagequeue.QueuedByUser, false, nil, + ) + if err != nil { + t.Fatalf("queue replacement: %v", err) + } + reserved, ok := svc.messageQueue.ReserveQueued(ctx, "session-1") + if !ok { + t.Fatal("reserve pending FIFO handoff") + } + reservation := svc.markQueuedDispatchInFlightWithSource("session-1", reserved.ID, reserved) + + claim, err := svc.sendNowRestoreClaimForReservation(ctx, reservation) + if err != nil { + t.Fatalf("build pending Send Now restore claim: %v", err) + } + if claim.SessionGeneration == 0 { + t.Fatal("pending Send Now restore claim omitted the session generation") + } + if err := svc.messageQueue.RestoreSendNowClaim(ctx, claim); err != nil { + t.Fatalf("restore pending Send Now claim: %v", err) + } + status := svc.messageQueue.GetStatus(ctx, "session-1") + if len(status.Entries) != 1 || status.Entries[0].ID != source.ID { + t.Fatalf("restored queue = %#v, want replacement entry", status.Entries) + } +} + +func TestPendingSendNowRestoreUsesReservationGeneration(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "task-1", "session-1", "step-1") + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + if _, err := svc.messageQueue.QueueMessage( + ctx, "session-1", "task-1", "pending", "", messagequeue.QueuedByUser, false, nil, + ); err != nil { + t.Fatalf("queue message: %v", err) + } + reserved, ok := svc.messageQueue.ReserveQueued(ctx, "session-1") + if !ok { + t.Fatal("reserve pending FIFO handoff") + } + reservation := svc.markQueuedDispatchInFlightWithSource("session-1", reserved.ID, reserved) + if _, err := svc.messageQueue.PurgeTask(ctx, "task-1"); err != nil { + t.Fatalf("purge task: %v", err) + } + + claim, err := svc.sendNowRestoreClaimForReservation(ctx, reservation) + if err != nil { + t.Fatalf("build pending Send Now restore claim: %v", err) + } + if claim.SessionGeneration != 0 { + t.Fatalf("restore session generation = %d, want reservation generation 0", claim.SessionGeneration) + } + if claim.SourceGenerations["task-1"] != 0 { + t.Fatalf("restore task generation = %d, want reservation generation 0", claim.SourceGenerations["task-1"]) + } +} func TestSendQueuedNowSupersedesPendingFIFOHandoff(t *testing.T) { ctx := context.Background() diff --git a/apps/backend/internal/orchestrator/queue_session_delete_notification_test.go b/apps/backend/internal/orchestrator/queue_session_delete_notification_test.go new file mode 100644 index 0000000000..40337b2eee --- /dev/null +++ b/apps/backend/internal/orchestrator/queue_session_delete_notification_test.go @@ -0,0 +1,171 @@ +package orchestrator + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/kandev/kandev/internal/events" + "github.com/kandev/kandev/internal/events/bus" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" + "github.com/kandev/kandev/internal/task/models" +) + +type blockingSessionAttachmentCleaner struct { + started chan struct{} + release chan struct{} +} + +func (c *blockingSessionAttachmentCleaner) DeleteSessionMessageAttachments(context.Context, string, string) error { + close(c.started) + <-c.release + return nil +} + +func TestDeletedSessionCleanupSerializesAttachmentCleanupWithQueueMutation(t *testing.T) { + ctx := context.Background() + queue := messagequeue.NewServiceMemory(testLogger()) + entry, err := queue.QueueMessage(ctx, "session-delete-race", "task-delete-race", "queued", "", messagequeue.QueuedByUser, false, nil) + if err != nil { + t.Fatalf("queue message: %v", err) + } + cleaner := &blockingSessionAttachmentCleaner{ + started: make(chan struct{}), + release: make(chan struct{}), + } + svc := &Service{ + messageQueue: queue, + sessionAttachmentCleaner: cleaner, + logger: testLogger(), + } + + cleanupDone := make(chan struct{}) + go func() { + svc.cancelDeletedSessionQueue(ctx, entry.TaskID, entry.SessionID) + close(cleanupDone) + }() + select { + case <-cleaner.started: + case <-time.After(time.Second): + t.Fatal("attachment cleanup did not start") + } + + mutationDone := make(chan error, 1) + go func() { + mutationDone <- queue.UpdateMessageWithMetadata( + ctx, entry.SessionID, entry.ID, "changed", nil, nil, messagequeue.QueuedByUser, + ) + }() + select { + case err := <-mutationDone: + t.Fatalf("queue mutation completed during attachment cleanup: %v", err) + case <-time.After(50 * time.Millisecond): + } + close(cleaner.release) + + select { + case <-cleanupDone: + case <-time.After(time.Second): + t.Fatal("session cleanup did not complete") + } + if err := <-mutationDone; err == nil { + t.Fatal("queue mutation unexpectedly succeeded after cleanup") + } +} + +func TestDeleteSessionPublishesOneQueueStatusNotification(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedTaskAndSession(t, repo, "task-session-delete-once", "session-delete-once", models.TaskSessionStateCompleted) + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + eventBus := bus.NewMemoryEventBus(testLogger()) + t.Cleanup(func() { eventBus.Close() }) + svc.eventBus = eventBus + repo.SetTaskSessionQueuePurgeNotifier(func(ctx context.Context, taskID, sessionID string) { + svc.purgeDeletedSessionQueue(ctx, taskID, sessionID) + }) + svc.sessionQueuePurgeNotifierRegistered = true + + var statusEvents atomic.Int32 + if _, err := eventBus.Subscribe(events.MessageQueueStatusChanged, func(context.Context, *bus.Event) error { + statusEvents.Add(1) + return nil + }); err != nil { + t.Fatalf("subscribe queue status: %v", err) + } + + if err := svc.DeleteSession(ctx, "session-delete-once"); err != nil { + t.Fatalf("delete session: %v", err) + } + if got := statusEvents.Load(); got != 1 { + t.Fatalf("queue status notifications = %d, want 1", got) + } +} + +func TestDeletedSessionQueueNotificationIgnoresCancelledContext(t *testing.T) { + queue := messagequeue.NewServiceMemory(testLogger()) + entry, err := queue.QueueMessage( + context.Background(), + "session-delete-cancelled", + "task-delete-cancelled", + "queued", + "", + messagequeue.QueuedByUser, + false, + nil, + ) + if err != nil { + t.Fatalf("queue message: %v", err) + } + + svc := &Service{messageQueue: queue, logger: testLogger()} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + svc.purgeDeletedSessionQueue(ctx, entry.TaskID, entry.SessionID) + + if got := queue.GetStatus(context.Background(), entry.SessionID).Count; got != 1 { + t.Fatalf("post-commit notification changed queue count = %d, want 1", got) + } +} + +func TestFallbackDeletedSessionCleanupPublishesQueueStatus(t *testing.T) { + queue := messagequeue.NewServiceMemory(testLogger()) + entry, err := queue.QueueMessage( + context.Background(), + "session-delete-fallback", + "task-delete-fallback", + "queued", + "", + messagequeue.QueuedByUser, + false, + nil, + ) + if err != nil { + t.Fatalf("queue message: %v", err) + } + eventBus := bus.NewMemoryEventBus(testLogger()) + t.Cleanup(func() { eventBus.Close() }) + svc := &Service{ + messageQueue: queue, + eventBus: eventBus, + logger: testLogger(), + } + var statusEvents atomic.Int32 + if _, err := eventBus.Subscribe(events.MessageQueueStatusChanged, func(_ context.Context, event *bus.Event) error { + data, _ := event.Data.(map[string]interface{}) + if data["session_id"] == entry.SessionID { + statusEvents.Add(1) + } + return nil + }); err != nil { + t.Fatalf("subscribe queue status: %v", err) + } + + svc.cancelDeletedSessionQueue(context.Background(), entry.TaskID, entry.SessionID) + + if got := statusEvents.Load(); got != 1 { + t.Fatalf("fallback queue status notifications = %d, want 1", got) + } +} diff --git a/apps/backend/internal/orchestrator/service.go b/apps/backend/internal/orchestrator/service.go index 31b8aa4b4b..05912ce122 100644 --- a/apps/backend/internal/orchestrator/service.go +++ b/apps/backend/internal/orchestrator/service.go @@ -54,6 +54,8 @@ var ( ErrRouteActionActiveTurn = errors.New("route actions require a settled turn") ) +const maxStartupTransferReconcileAttempts = 30 + // ServiceConfig holds orchestrator service configuration type ServiceConfig struct { Scheduler scheduler.SchedulerConfig @@ -147,6 +149,23 @@ type TransientRetryMessageService interface { DeleteMessage(ctx context.Context, id string) error } +// SessionAttachmentCleaner removes file-backed prompt attachments after a +// task session has been deleted. The task service owns storage and authorization +// details; the orchestrator only coordinates the post-commit lifecycle hook. +type SessionAttachmentCleaner interface { + DeleteSessionMessageAttachments(ctx context.Context, taskID, sessionID string) error +} + +// SessionAttachmentTransferer keeps the exact claimed prompt attachments +// represented by a queue transfer bound to the queue's owning session. +type SessionAttachmentTransferer interface { + TransferSessionMessageAttachments( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + attachmentIDs []string, + ) error +} + // SubagentContextRecorder persists a durable relational record of a subagent // (Task tool) invocation observed on a tool-call frame. It returns nothing — // a repository failure never fails the enclosing message write, turn, or @@ -615,7 +634,15 @@ type Service struct { // transientRetryMessages owns durable cleanup of persisted retry notices. // It is optional for focused tests and pre-composition callers. transientRetryMessages TransientRetryMessageService + // sessionQueuePurgeNotifierRegistered means the task repository owns + // queue cleanup and status publication after DeleteTaskSession commits. + // DeleteSession keeps its fallback cleanup for focused compositions that + // cannot register the repository callback, but must not publish twice when + // the callback is active. + sessionQueuePurgeNotifierRegistered bool + sessionAttachmentCleaner SessionAttachmentCleaner + sessionAttachmentTransferer SessionAttachmentTransferer // subagentContexts optionally persists a relational record of subagent // (Task tool) invocations recognized on the tool-call frame paths. Nil is // safe: both call sites guard on it. See SetSubagentContextRecorder. @@ -1260,14 +1287,15 @@ type Service struct { running bool startedAt time.Time - // sendNowWorkers owns the asynchronous replacement handoffs. The context - // is cancelled during Stop so a claimed durable source gets a bounded - // recovery attempt before shutdown returns. + // sendNowWorkers owns one generation of asynchronous replacement handoffs. + // Stop cancels the generation and waits only for a bounded interval; a late + // generation keeps its own wait group while the next Start creates another. sendNowMu sync.Mutex sendNowCtx context.Context sendNowCancel context.CancelFunc sendNowStopped bool - sendNowWorkers sync.WaitGroup + sendNowWorkers *sync.WaitGroup + sendNowDrain chan struct{} // ciAutomationWorkers owns the asynchronous per-PR automation loops. The // service-owned context lets Stop cancel in-flight evaluations and prevents @@ -1525,6 +1553,7 @@ func NewService( reservedPromptCallbacks: newReservedPromptCallbackOwner(), sendNowCtx: sendNowCtx, sendNowCancel: sendNowCancel, + sendNowWorkers: &sync.WaitGroup{}, ciAutomationCtx: ciAutomationCtx, ciAutomationCancel: ciAutomationCancel, dynamicSuccessorCtx: dynamicSuccessorCtx, @@ -1552,9 +1581,21 @@ func NewService( SetTaskQueuePurgeNotifier(func(context.Context, string)) }); ok { registrar.SetTaskQueuePurgeNotifier(func(ctx context.Context, taskID string) { + // SQLite task lifecycle purges the durable queue in its own + // transaction, so invalidate only the process-local edit leases + // here. Calling PurgeTask would race with the next queue generation. + msgQueue.InvalidateEditLeasesForTask(taskID) s.publishTaskQueueStatusEvent(ctx, taskID, "") }) } + if registrar, ok := repo.(interface { + SetTaskSessionQueuePurgeNotifier(func(context.Context, string, string)) + }); ok { + s.sessionQueuePurgeNotifierRegistered = true + registrar.SetTaskSessionQueuePurgeNotifier(func(ctx context.Context, taskID, sessionID string) { + s.purgeDeletedSessionQueue(ctx, taskID, sessionID) + }) + } s.backgroundProbe = serviceBackgroundProbeAdapter{s: s} exec.SetOnContextWindowReset(s.clearContextWindowForReset) @@ -1764,6 +1805,18 @@ func (s *Service) SetLaunchAttachmentClaimer(claimer LaunchAttachmentClaimer) { s.launchAttachmentClaimer = claimer } +// SetSessionAttachmentCleaner wires post-delete cleanup for claimed prompt +// attachments belonging to a task session. +func (s *Service) SetSessionAttachmentCleaner(cleaner SessionAttachmentCleaner) { + s.sessionAttachmentCleaner = cleaner +} + +// SetSessionAttachmentTransferer wires the attachment registry update required +// when queued work moves between task sessions. +func (s *Service) SetSessionAttachmentTransferer(transfer SessionAttachmentTransferer) { + s.sessionAttachmentTransferer = transfer +} + // SetOnPrimarySessionSet sets a callback on the executor for when the first session // of a task is marked primary. Used to publish a task.updated event so the frontend // receives the primary_session_id. @@ -2793,6 +2846,39 @@ func (s *Service) isSessionResetInProgress(sessionID string) bool { return inProgress } +func (s *Service) reconcileDurableQueueStateOnStartup(ctx context.Context) error { + for attempt := 1; ; attempt++ { + err := s.reconcileSessionTransferCompensationsOnStartup(ctx) + if !errors.Is(err, messagequeue.ErrSessionTransferInProgress) { + if err != nil { + return fmt.Errorf("reconcile session transfer compensations: %w", err) + } + break + } + if attempt == maxStartupTransferReconcileAttempts { + return fmt.Errorf( + "reconcile session transfer compensations still blocked after %d attempts: %w", + attempt, + err, + ) + } + timer := time.NewTimer(time.Second) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case <-timer.C: + } + } + if err := s.reconcilePendingQueueDispatchesOnStartup(ctx); err != nil { + return fmt.Errorf("reconcile pending queue dispatches: %w", err) + } + if err := s.reconcilePendingSendNowClaimsOnStartup(ctx); err != nil { + return fmt.Errorf("reconcile pending Send Now claims: %w", err) + } + return nil +} + // Start starts all orchestrator components func (s *Service) Start(ctx context.Context) error { s.mu.Lock() @@ -2812,11 +2898,22 @@ func (s *Service) Start(ctx context.Context) error { s.logger.Info("starting orchestrator service") s.resetReservedPromptCallbacks() - s.resetSendNowWorkers() + if err := s.resetSendNowWorkers(); err != nil { + s.mu.Lock() + s.running = false + s.mu.Unlock() + return err + } s.resetCIAutomationWorkers() s.resetDynamicSuccessorWorkers() + if err := s.reconcileDurableQueueStateOnStartup(ctx); err != nil { + s.logger.Error("failed to reconcile durable queue state on startup", zap.Error(err)) + s.mu.Lock() + s.running = false + s.mu.Unlock() + return err + } s.resetParkedSamplingWorkers() - // Reconcile session state from persisted runtime state on startup. // This does NOT launch any agent processes — sessions are recovered lazily // when the user opens them (via task.session.status → task.session.resume). diff --git a/apps/backend/internal/orchestrator/task_operations.go b/apps/backend/internal/orchestrator/task_operations.go index 7316550410..85a6fd50a7 100644 --- a/apps/backend/internal/orchestrator/task_operations.go +++ b/apps/backend/internal/orchestrator/task_operations.go @@ -3678,8 +3678,11 @@ func (s *Service) DeleteSession(ctx context.Context, sessionID string) error { s.clearTurnActivity(sessionID) // Queue rows and per-session policy are removed atomically with the task - // session row. Publish only the task-scoped recount after commit. - s.publishTaskQueueStatusEvent(ctx, taskID, "") + // session row. The repository callback publishes a session-scoped event + // when it owns the purge; focused compositions use the task-scoped fallback. + if !s.sessionQueuePurgeNotifierRegistered { + s.publishTaskQueueStatusEvent(ctx, taskID, "") + } // Auto-promote another session if we deleted the primary if wasPrimary { @@ -3818,6 +3821,60 @@ func (s *Service) publishTaskSessionErrorEvent( )) } +// purgeDeletedSessionQueue invalidates in-process edit state and publishes the +// queue update after the task repository commits the durable session purge. +// Request cancellation must not prevent this post-commit notification. +func (s *Service) purgeDeletedSessionQueue(ctx context.Context, taskID, sessionID string) { + cleanupCtx := context.WithoutCancel(ctx) + if s.messageQueue == nil { + return + } + s.messageQueue.InvalidateEditLeasesForSession(sessionID) + s.publishTaskQueueStatusEvent(cleanupCtx, taskID, sessionID) +} + +// cancelDeletedSessionQueue removes file-backed prompt attachments left on a +// deleted session. Queue rows and their status notification are owned by the +// repository's post-commit callback when that callback is registered. Focused +// compositions without the callback retain the direct queue purge fallback. +func (s *Service) cancelDeletedSessionQueue(ctx context.Context, taskID, sessionID string) { + cleanupCtx := context.WithoutCancel(ctx) + if s.messageQueue == nil { + if s.sessionAttachmentCleaner != nil { + if err := s.sessionAttachmentCleaner.DeleteSessionMessageAttachments(cleanupCtx, taskID, sessionID); err != nil { + s.logger.Warn("failed to remove session attachment bytes after session delete", + zap.String("session_id", sessionID), + zap.String("task_id", taskID), + zap.Error(err)) + } + } + return + } + _ = s.messageQueue.WithSessionAdmission(cleanupCtx, sessionID, func(admittedCtx context.Context) error { + if s.sessionAttachmentCleaner != nil { + if err := s.sessionAttachmentCleaner.DeleteSessionMessageAttachments(admittedCtx, taskID, sessionID); err != nil { + s.logger.Warn("failed to remove session attachment bytes after session delete", + zap.String("session_id", sessionID), + zap.String("task_id", taskID), + zap.Error(err)) + } + } + if s.sessionQueuePurgeNotifierRegistered { + return nil + } + if _, err := s.messageQueue.PurgeSession(admittedCtx, sessionID); err != nil { + s.logger.Warn("failed to purge queued prompts after session delete", + zap.String("session_id", sessionID), + zap.String("task_id", taskID), + zap.Error(err)) + } + return nil + }) + if !s.sessionQueuePurgeNotifierRegistered { + s.publishTaskQueueStatusEvent(cleanupCtx, taskID, sessionID) + } +} + // quiesceSessionExecutionBeforeDeletion stops the in-memory lifecycle // execution, if one exists, before a session row is removed. A runtime that // explicitly reports the exact execution as absent is already quiesced; other @@ -4385,6 +4442,7 @@ type promptTaskOptions struct { claimEntryID string lifecyclePrompt bool afterClaim func() error + afterDispatch func() error preservePromptContext bool // reserveTurnUntilDispatch persists detached-resume ownership before agentctl // dispatch, while delaying the visible turn.started event until acceptance. @@ -4665,6 +4723,15 @@ func (s *Service) promptTask(ctx context.Context, taskID, sessionID string, prom promptCtx, taskID, sessionID, rollback.sessionIdentity, session.AgentExecutionID, rollback.reservedTurn, foregroundDispatch, dispatchOutcome, ) + if options.afterDispatch != nil { + originalOnDispatched := onDispatched + onDispatched = func() { + durableOutcomeErr := options.afterDispatch() + originalOnDispatched() + _, publicationErr := dispatchOutcome.snapshot() + dispatchOutcome.recordAccepted(errors.Join(durableOutcomeErr, publicationErr)) + } + } if releaseDispatchGuard != nil { originalOnDispatched := onDispatched onDispatched = func() { @@ -5974,6 +6041,40 @@ func (s *Service) DrainQueuedMessage(ctx context.Context, sessionID string) (boo return s.drainQueuedMessageForPromptableSessionLocked(ctx, sessionID), nil } +// DrainQueuedMessageIfAutoRun dispatches one queued message only when the +// persisted Auto-run policy is already enabled. Unlike DrainQueuedMessage, it +// never changes the policy as a side effect. +func (s *Service) DrainQueuedMessageIfAutoRun(ctx context.Context, sessionID string) (bool, error) { + if sessionID == "" { + return false, fmt.Errorf("session_id is required") + } + lock, release := s.acquireCancelInFlightGuard(sessionID) + defer release() + lock.Lock() + defer lock.Unlock() + if s.isCancelInFlight(sessionID) || s.isQueuedDispatchInFlight(sessionID) || s.isSteerInFlight(sessionID) { + return false, nil + } + if s.sessionHasPendingClarification(ctx, sessionID) { + return false, nil + } + if s.messageQueue == nil { + return false, errors.New("message queue is not configured") + } + status := s.messageQueue.GetStatus(ctx, sessionID) + if !status.AutoRun { + return false, nil + } + session, err := s.repo.GetTaskSession(ctx, sessionID) + if err != nil { + return false, fmt.Errorf("failed to get session: %w", err) + } + if err := s.checkSessionPromptable(session.TaskID, sessionID, session.State); err != nil { + return false, nil + } + return s.drainQueuedMessageForPromptableSessionLocked(ctx, sessionID), nil +} + // DrainQueuedMessageForSession dispatches one message only for the exact session incarnation. func (s *Service) DrainQueuedMessageForSession(ctx context.Context, identity messagequeue.QueueSessionIdentity) (bool, error) { if identity.TaskID == "" || identity.SessionID == "" || identity.SessionIncarnationID == "" { diff --git a/apps/backend/internal/orchestrator/task_operations_test.go b/apps/backend/internal/orchestrator/task_operations_test.go index cdf8a01708..b4c5c5c21a 100644 --- a/apps/backend/internal/orchestrator/task_operations_test.go +++ b/apps/backend/internal/orchestrator/task_operations_test.go @@ -2187,6 +2187,25 @@ func TestCancelAgent_QueuedMessageRunsAfterExplicitDrain(t *testing.T) { } } +func TestDrainQueuedMessageIfAutoRunDoesNotResumePausedQueue(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + svc := createTestServiceWithAgent(repo, newMockStepGetter(), newMockTaskRepo(), &mockAgentManager{}) + seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateWaitingForInput) + _, err := svc.messageQueue.QueueMessage( + ctx, "session1", "task1", "paused queued message", "", messagequeue.QueuedByUser, false, nil, + ) + require.NoError(t, err) + require.NoError(t, svc.messageQueue.SetAutoRun(ctx, "session1", false)) + + drained, err := svc.DrainQueuedMessageIfAutoRun(ctx, "session1") + require.NoError(t, err) + require.False(t, drained) + status := svc.messageQueue.GetStatus(ctx, "session1") + require.False(t, status.AutoRun) + require.Len(t, status.Entries, 1) +} + func queueAndInterruptForPeerMessage( svc *Service, ctx context.Context, diff --git a/apps/backend/internal/task/repository/interface.go b/apps/backend/internal/task/repository/interface.go index 4302154bb9..9f1c809db9 100644 --- a/apps/backend/internal/task/repository/interface.go +++ b/apps/backend/internal/task/repository/interface.go @@ -280,6 +280,8 @@ type AttachmentRepository interface { ClaimMessageAttachments(ctx context.Context, ids []string, ownerID, workspaceID, taskID, sessionID string) error DeleteClaimedMessageAttachments(ctx context.Context, ids []string, ownerID, taskID, sessionID string) ([]*models.TaskMessageAttachment, error) DeleteMessageAttachmentsByTask(ctx context.Context, taskID string) ([]*models.TaskMessageAttachment, error) + DeleteMessageAttachmentsBySession(ctx context.Context, taskID, sessionID string) ([]*models.TaskMessageAttachment, error) + TransferMessageAttachments(ctx context.Context, taskID, oldSessionID, newSessionID string, attachmentIDs []string) error DeleteMessageAttachment(ctx context.Context, id, ownerID string) error MarkExpiredMessageAttachments(ctx context.Context, now time.Time) ([]*models.TaskMessageAttachment, error) } diff --git a/apps/backend/internal/task/repository/sqlite/attachment.go b/apps/backend/internal/task/repository/sqlite/attachment.go index 9f6fef2429..99e93fd62c 100644 --- a/apps/backend/internal/task/repository/sqlite/attachment.go +++ b/apps/backend/internal/task/repository/sqlite/attachment.go @@ -11,6 +11,7 @@ import ( "github.com/google/uuid" "github.com/jmoiron/sqlx" "github.com/kandev/kandev/internal/db/dialect" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" "github.com/kandev/kandev/internal/task/models" ) @@ -102,11 +103,20 @@ func (r *Repository) ClaimMessageAttachments(ctx context.Context, ids []string, if len(ids) > models.MaxMessageAttachmentCount { return models.ErrTooManyAttachments } + queueLockPresent, err := r.queueSessionLockTablePresent(ctx) + if err != nil { + return fmt.Errorf("check queue session lock schema: %w", err) + } tx, err := r.db.BeginTxx(ctx, nil) if err != nil { return fmt.Errorf("begin attachment claim: %w", err) } defer func() { _ = tx.Rollback() }() + if queueLockPresent { + if err := messagequeue.LockSessionInTransaction(ctx, tx, r.db, sessionID); err != nil { + return fmt.Errorf("lock attachment claim session: %w", err) + } + } if taskID != "" { result, err := tx.ExecContext(ctx, tx.Rebind(` UPDATE tasks SET updated_at = updated_at WHERE id = ? @@ -233,7 +243,25 @@ func (r *Repository) DeleteMessageAttachment(ctx context.Context, id, ownerID st return nil } -func (r *Repository) DeleteClaimedMessageAttachments(ctx context.Context, ids []string, ownerID, taskID, sessionID string) ([]*models.TaskMessageAttachment, error) { +func (r *Repository) DeleteClaimedMessageAttachments( + ctx context.Context, ids []string, ownerID, taskID, sessionID string, +) ([]*models.TaskMessageAttachment, error) { + return r.deleteClaimedMessageAttachments(ctx, ids, ownerID, taskID, sessionID, true) +} + +func (r *Repository) DeleteClaimedMessageAttachmentsByTaskSession( + ctx context.Context, ids []string, taskID, sessionID string, +) ([]*models.TaskMessageAttachment, error) { + return r.deleteClaimedMessageAttachments(ctx, ids, "", taskID, sessionID, false) +} + +//nolint:cyclop,funlen // Attachment release coordinates ownership, queue locks, and transactional cleanup. +func (r *Repository) deleteClaimedMessageAttachments( + ctx context.Context, + ids []string, + ownerID, taskID, sessionID string, + ownerRequired bool, +) ([]*models.TaskMessageAttachment, error) { if len(ids) == 0 { return nil, nil } @@ -241,11 +269,21 @@ func (r *Repository) DeleteClaimedMessageAttachments(ctx context.Context, ids [] if err != nil { return nil, fmt.Errorf("probe queue attachment references: %w", err) } + queueLockPresent := false + if sessionID != "" { + queueLockPresent, err = r.queueSessionLockTablePresent(ctx) + if err != nil { + return nil, fmt.Errorf("check queue session lock schema: %w", err) + } + } tx, err := r.db.BeginTxx(ctx, nil) if err != nil { return nil, fmt.Errorf("begin claimed attachment release: %w", err) } defer func() { _ = tx.Rollback() }() + if err := r.lockClaimedAttachmentReleaseTx(ctx, tx, sessionID, queueLockPresent); err != nil { + return nil, err + } if sessionID != "" { if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), sessionID); err != nil { return nil, err @@ -279,47 +317,24 @@ func (r *Repository) DeleteClaimedMessageAttachments(ctx context.Context, ids [] return nil, nil } placeholders := strings.TrimSuffix(strings.Repeat("?,", len(releasable)), ",") - args := []interface{}{ownerID, taskID, sessionID, models.AttachmentStateClaimed} + ownerClause := "" + args := []interface{}{taskID, sessionID, models.AttachmentStateClaimed} + if ownerRequired { + ownerClause = "owner_id = ? AND " + args = append([]interface{}{ownerID}, args...) + } args = append(args, idsToInterfaces(releasable)...) - rows, err := tx.QueryxContext(ctx, tx.Rebind(` - SELECT `+attachmentSelectColumns+` FROM task_message_attachments - WHERE owner_id = ? AND task_id = ? AND session_id = ? AND state = ? - AND id IN (`+placeholders+`) - `), args...) + released, err := listClaimedAttachmentsForReleaseTx( + ctx, tx, ownerClause, attachmentSessionScopeClause(sessionID), placeholders, args, + ) if err != nil { - return nil, fmt.Errorf("list claimed attachments for release: %w", err) + return nil, err } - var released []*models.TaskMessageAttachment - for rows.Next() { - attachment := &models.TaskMessageAttachment{} - if err := rows.StructScan(attachment); err != nil { - _ = rows.Close() - return nil, fmt.Errorf("scan claimed attachment for release: %w", err) - } - released = append(released, attachment) - } - if err := rows.Err(); err != nil { - _ = rows.Close() - return nil, fmt.Errorf("iterate claimed attachments for release: %w", err) - } - if err := rows.Close(); err != nil { - return nil, fmt.Errorf("close claimed attachments for release: %w", err) - } - for _, attachment := range released { - result, err := tx.ExecContext(ctx, tx.Rebind(` - DELETE FROM task_message_attachments - WHERE id = ? AND owner_id = ? AND task_id = ? AND session_id = ? AND state = ? - `), attachment.ID, ownerID, taskID, sessionID, models.AttachmentStateClaimed) - if err != nil { - return nil, fmt.Errorf("release claimed attachment: %w", err) - } - affected, err := result.RowsAffected() - if err != nil { - return nil, fmt.Errorf("release claimed attachment rows affected: %w", err) - } - if affected != 1 { - return nil, models.ErrAttachmentClaimConflict - } + if err := deleteClaimedAttachmentsTx( + ctx, tx, released, ownerClause, attachmentSessionScopeClause(sessionID), + ownerID, ownerRequired, taskID, sessionID, + ); err != nil { + return nil, err } if err := tx.Commit(); err != nil { return nil, fmt.Errorf("commit claimed attachment release: %w", err) @@ -327,6 +342,18 @@ func (r *Repository) DeleteClaimedMessageAttachments(ctx context.Context, ids [] return released, nil } +func (r *Repository) lockClaimedAttachmentReleaseTx( + ctx context.Context, tx *sqlx.Tx, sessionID string, queueLockPresent bool, +) error { + if !queueLockPresent { + return nil + } + if err := messagequeue.LockSessionInTransaction(ctx, tx, r.db, sessionID); err != nil { + return fmt.Errorf("lock attachment release session %s: %w", sessionID, err) + } + return nil +} + func lockAttachmentQueueSession( ctx context.Context, tx *sqlx.Tx, @@ -353,6 +380,77 @@ func lockAttachmentQueueSession( return nil } +func listClaimedAttachmentsForReleaseTx( + ctx context.Context, + tx *sqlx.Tx, + ownerClause, sessionScopeClause, placeholders string, + args []interface{}, +) ([]*models.TaskMessageAttachment, error) { + rows, err := tx.QueryxContext(ctx, tx.Rebind(` + SELECT `+attachmentSelectColumns+` FROM task_message_attachments + WHERE `+ownerClause+`task_id = ? AND `+sessionScopeClause+` AND state = ? + AND id IN (`+placeholders+`) + `), args...) + if err != nil { + return nil, fmt.Errorf("list claimed attachments for release: %w", err) + } + defer func() { _ = rows.Close() }() + var released []*models.TaskMessageAttachment + for rows.Next() { + attachment := &models.TaskMessageAttachment{} + if err := rows.StructScan(attachment); err != nil { + return nil, fmt.Errorf("scan claimed attachment for release: %w", err) + } + released = append(released, attachment) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate claimed attachments for release: %w", err) + } + if err := rows.Close(); err != nil { + return nil, fmt.Errorf("close claimed attachments for release: %w", err) + } + return released, nil +} + +func deleteClaimedAttachmentsTx( + ctx context.Context, + tx *sqlx.Tx, + released []*models.TaskMessageAttachment, + ownerClause, sessionScopeClause, ownerID string, + ownerRequired bool, + taskID, sessionID string, +) error { + for _, attachment := range released { + deleteArgs := []interface{}{attachment.ID} + if ownerRequired { + deleteArgs = append(deleteArgs, ownerID) + } + deleteArgs = append(deleteArgs, taskID, sessionID, models.AttachmentStateClaimed) + result, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM task_message_attachments + WHERE id = ? AND `+ownerClause+`task_id = ? AND `+sessionScopeClause+` AND state = ? + `), deleteArgs...) + if err != nil { + return fmt.Errorf("release claimed attachment: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("release claimed attachment rows affected: %w", err) + } + if affected != 1 { + return models.ErrAttachmentClaimConflict + } + } + return nil +} + +func attachmentSessionScopeClause(sessionID string) string { + if sessionID == "" { + return "session_id = ?" + } + return "(session_id = ? OR session_id = '')" +} + type attachmentReference struct { AttachmentID string `json:"attachment_id"` } @@ -609,6 +707,52 @@ func (r *Repository) DeleteMessageAttachmentsByTask(ctx context.Context, taskID return attachments, nil } +// DeleteMessageAttachmentsBySession removes claimed attachment descriptors +// owned by a deleted task session. Staged uploads have no session binding and +// remain governed by their expiry cleanup. +func (r *Repository) DeleteMessageAttachmentsBySession(ctx context.Context, taskID, sessionID string) ([]*models.TaskMessageAttachment, error) { + if taskID == "" || sessionID == "" { + return nil, nil + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin session attachment cleanup: %w", err) + } + defer func() { _ = tx.Rollback() }() + rows, err := tx.QueryxContext(ctx, tx.Rebind(` + SELECT `+attachmentSelectColumns+` FROM task_message_attachments + WHERE task_id = ? AND session_id = ? + `), taskID, sessionID) + if err != nil { + return nil, fmt.Errorf("list session attachments for cleanup: %w", err) + } + var attachments []*models.TaskMessageAttachment + for rows.Next() { + attachment := &models.TaskMessageAttachment{} + if err := rows.StructScan(attachment); err != nil { + _ = rows.Close() + return nil, fmt.Errorf("scan session attachment for cleanup: %w", err) + } + attachments = append(attachments, attachment) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return nil, fmt.Errorf("iterate session attachments for cleanup: %w", err) + } + if err := rows.Close(); err != nil { + return nil, fmt.Errorf("close session attachments for cleanup: %w", err) + } + if _, err := tx.ExecContext(ctx, tx.Rebind(` + DELETE FROM task_message_attachments WHERE task_id = ? AND session_id = ? + `), taskID, sessionID); err != nil { + return nil, fmt.Errorf("delete session attachments: %w", err) + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit session attachment cleanup: %w", err) + } + return attachments, nil +} + func (r *Repository) DeleteMessageAttachmentsByWorkspaceTx( ctx context.Context, tx *sqlx.Tx, @@ -623,28 +767,169 @@ func (r *Repository) DeleteMessageAttachmentsByWorkspaceTx( if err != nil { return nil, fmt.Errorf("list workspace attachments for cleanup: %w", err) } + defer func() { _ = rows.Close() }() var attachments []*models.TaskMessageAttachment for rows.Next() { attachment := &models.TaskMessageAttachment{} if err := rows.StructScan(attachment); err != nil { - _ = rows.Close() return nil, fmt.Errorf("scan workspace attachment for cleanup: %w", err) } attachments = append(attachments, attachment) } if err := rows.Err(); err != nil { - _ = rows.Close() return nil, fmt.Errorf("iterate workspace attachments for cleanup: %w", err) } - if err := rows.Close(); err != nil { - return nil, fmt.Errorf("close workspace attachments for cleanup: %w", err) - } if _, err := tx.ExecContext(ctx, tx.Rebind(`DELETE FROM task_message_attachments WHERE workspace_id = ?`), workspaceID); err != nil { return nil, fmt.Errorf("delete workspace attachments: %w", err) } return attachments, nil } +func (r *Repository) TransferMessageAttachments( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + attachmentIDs []string, +) error { + if taskID == "" || oldSessionID == "" || newSessionID == "" || + oldSessionID == newSessionID || len(attachmentIDs) == 0 { + return nil + } + uniqueAttachmentIDs := uniqueAttachmentIDs(attachmentIDs) + queueLockPresent, err := r.queueSessionLockTablePresent(ctx) + if err != nil { + return fmt.Errorf("check queue session lock schema: %w", err) + } + tx, err := r.db.BeginTxx(ctx, nil) + if err != nil { + return fmt.Errorf("begin transfer session attachments: %w", err) + } + defer func() { _ = tx.Rollback() }() + if queueLockPresent { + if err := messagequeue.LockSessionPairInTransaction( + ctx, tx, r.db, oldSessionID, newSessionID, + ); err != nil { + return fmt.Errorf("lock transfer session pair: %w", err) + } + } + locations, err := transferAttachmentLocations(ctx, tx, taskID, uniqueAttachmentIDs) + if err != nil { + return err + } + sourceAttachmentIDs, err := validateTransferAttachmentLocations( + taskID, oldSessionID, newSessionID, uniqueAttachmentIDs, locations, + ) + if err != nil { + return err + } + if err := updateTransferSessionAttachments( + ctx, tx, taskID, oldSessionID, newSessionID, sourceAttachmentIDs, + ); err != nil { + return err + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit transfer session attachments: %w", err) + } + return nil +} + +func uniqueAttachmentIDs(attachmentIDs []string) []string { + unique := make([]string, 0, len(attachmentIDs)) + seen := make(map[string]struct{}, len(attachmentIDs)) + for _, attachmentID := range attachmentIDs { + if _, ok := seen[attachmentID]; ok { + continue + } + seen[attachmentID] = struct{}{} + unique = append(unique, attachmentID) + } + return unique +} + +func transferAttachmentLocations( + ctx context.Context, tx *sqlx.Tx, taskID string, attachmentIDs []string, +) (map[string]string, error) { + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(attachmentIDs)), ",") + args := make([]interface{}, 0, 2+len(attachmentIDs)) + args = append(args, taskID, models.AttachmentStateClaimed) + for _, attachmentID := range attachmentIDs { + args = append(args, attachmentID) + } + rows, err := tx.QueryxContext(ctx, tx.Rebind(` + SELECT id, session_id + FROM task_message_attachments + WHERE task_id = ? AND state = ? AND id IN (`+placeholders+`) + `), args...) + if err != nil { + return nil, fmt.Errorf("inspect transfer session attachments: %w", err) + } + defer func() { _ = rows.Close() }() + locations := make(map[string]string, len(attachmentIDs)) + for rows.Next() { + var attachmentID, sessionID string + if err := rows.Scan(&attachmentID, &sessionID); err != nil { + return nil, fmt.Errorf("scan transfer session attachment: %w", err) + } + locations[attachmentID] = sessionID + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate transfer session attachments: %w", err) + } + if err := rows.Close(); err != nil { + return nil, fmt.Errorf("close transfer session attachments: %w", err) + } + return locations, nil +} + +func validateTransferAttachmentLocations( + taskID, oldSessionID, newSessionID string, attachmentIDs []string, locations map[string]string, +) ([]string, error) { + sourceAttachmentIDs := make([]string, 0, len(attachmentIDs)) + for _, attachmentID := range attachmentIDs { + sessionID, ok := locations[attachmentID] + if !ok || (sessionID != "" && sessionID != oldSessionID && sessionID != newSessionID) { + return nil, fmt.Errorf( + "transfer session attachments: attachment %q is not a claimed attachment for task %q in either transfer session", + attachmentID, taskID, + ) + } + if sessionID == oldSessionID { + sourceAttachmentIDs = append(sourceAttachmentIDs, attachmentID) + } + } + return sourceAttachmentIDs, nil +} + +func updateTransferSessionAttachments( + ctx context.Context, tx *sqlx.Tx, taskID, oldSessionID, newSessionID string, attachmentIDs []string, +) error { + if len(attachmentIDs) == 0 { + return nil + } + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(attachmentIDs)), ",") + args := make([]interface{}, 0, 5+len(attachmentIDs)) + args = append(args, newSessionID, time.Now().UTC(), taskID, oldSessionID, models.AttachmentStateClaimed) + for _, attachmentID := range attachmentIDs { + args = append(args, attachmentID) + } + result, err := tx.ExecContext(ctx, tx.Rebind(` + UPDATE task_message_attachments + SET session_id = ?, updated_at = ? + WHERE task_id = ? AND session_id = ? AND state = ? + AND id IN (`+placeholders+`) + `), args...) + if err != nil { + return fmt.Errorf("transfer session attachments: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("count transferred session attachments: %w", err) + } + if affected != int64(len(attachmentIDs)) { + return fmt.Errorf("transfer session attachments: expected %d claimed source attachments, moved %d", len(attachmentIDs), affected) + } + return nil +} + func (r *Repository) MarkExpiredMessageAttachments(ctx context.Context, now time.Time) ([]*models.TaskMessageAttachment, error) { tx, err := r.db.BeginTxx(ctx, nil) if err != nil { diff --git a/apps/backend/internal/task/repository/sqlite/attachment_postgres_test.go b/apps/backend/internal/task/repository/sqlite/attachment_postgres_test.go index 60af02118f..0db9d547e9 100644 --- a/apps/backend/internal/task/repository/sqlite/attachment_postgres_test.go +++ b/apps/backend/internal/task/repository/sqlite/attachment_postgres_test.go @@ -5,6 +5,7 @@ import ( "testing" "time" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" "github.com/kandev/kandev/internal/task/models" "github.com/kandev/kandev/internal/testutil" ) @@ -45,6 +46,80 @@ func TestPostgresMessageAttachmentLifecycle(t *testing.T) { t.Fatalf("released = %+v", released) } }) + t.Run("task/session cleanup without owner", func(t *testing.T) { + now := time.Now().UTC() + attachment := &models.TaskMessageAttachment{ + ID: "attachment-cleanup-session-pg", OwnerID: "owner-pg", WorkspaceID: "workspace-attachments-pg", + TaskID: "task-cleanup-session-pg", SessionID: "session-cleanup-pg", Name: "cleanup-session.bin", + MimeType: "application/octet-stream", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, + StorageKey: "attachment-cleanup-session-pg", State: models.AttachmentStateClaimed, + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + } + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatalf("create attachment: %v", err) + } + removed, err := repo.DeleteClaimedMessageAttachmentsByTaskSession( + ctx, []string{attachment.ID}, attachment.TaskID, attachment.SessionID, + ) + if err != nil { + t.Fatalf("cleanup release: %v", err) + } + if len(removed) != 1 || removed[0].ID != attachment.ID { + t.Fatalf("removed = %+v", removed) + } + }) + + t.Run("owned transfer during active compensation", func(t *testing.T) { + taskID := "task-transfer-pg" + oldSessionID := "session-transfer-old-pg" + newSessionID := "session-transfer-new-pg" + if err := repo.CreateTask(ctx, &models.Task{ID: taskID, Title: "Attachment transfer"}); err != nil { + t.Fatalf("create task: %v", err) + } + for _, sessionID := range []string{oldSessionID, newSessionID} { + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: sessionID, TaskID: taskID, State: models.TaskSessionStateWaitingForInput, + }); err != nil { + t.Fatalf("create session %s: %v", sessionID, err) + } + } + queueRepo, err := messagequeue.NewSQLiteRepository(db, db) + if err != nil { + t.Fatalf("create queue repository: %v", err) + } + compensations := queueRepo.(interface { + UpsertSessionTransferCompensation(context.Context, messagequeue.SessionTransferCompensation) error + }) + if err := compensations.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "attachment-transfer-pg", + TaskID: taskID, FromSessionID: oldSessionID, ToSessionID: newSessionID, + }); err != nil { + t.Fatalf("create compensation: %v", err) + } + now := time.Now().UTC() + attachment := &models.TaskMessageAttachment{ + ID: "attachment-transfer-pg", OwnerID: "owner-pg", WorkspaceID: "workspace-attachments-pg", + TaskID: taskID, SessionID: oldSessionID, Name: "transfer.bin", + MimeType: "application/octet-stream", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, + StorageKey: "attachment-transfer-pg", State: models.AttachmentStateClaimed, + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + } + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatalf("create transfer attachment: %v", err) + } + if err := repo.TransferMessageAttachments( + ctx, taskID, oldSessionID, newSessionID, []string{attachment.ID}, + ); err != nil { + t.Fatalf("transfer attachment: %v", err) + } + stored, err := repo.GetMessageAttachment(ctx, attachment.ID) + if err != nil { + t.Fatalf("load transferred attachment: %v", err) + } + if stored.SessionID != newSessionID { + t.Fatalf("transferred session = %q, want %q", stored.SessionID, newSessionID) + } + }) t.Run("task cleanup", func(t *testing.T) { now := time.Now().UTC() diff --git a/apps/backend/internal/task/repository/sqlite/attachment_test.go b/apps/backend/internal/task/repository/sqlite/attachment_test.go index b7bb8d87f5..ce0d11911a 100644 --- a/apps/backend/internal/task/repository/sqlite/attachment_test.go +++ b/apps/backend/internal/task/repository/sqlite/attachment_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + "github.com/kandev/kandev/internal/orchestrator/messagequeue" "github.com/kandev/kandev/internal/task/models" ) @@ -316,6 +317,70 @@ func TestDeleteClaimedMessageAttachments_RemovesOnlyMatchingClaims(t *testing.T) } } +func TestDeleteClaimedMessageAttachmentsByTaskSession_ReleasesWithoutOwner(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + now := time.Now().UTC() + attachments := []*models.TaskMessageAttachment{ + {ID: "cleanup-release", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-1", SessionID: "session-1", Name: "one", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "cleanup-release", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + {ID: "cleanup-release-other", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-2", SessionID: "session-1", Name: "two", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "cleanup-release-other", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + } + for _, attachment := range attachments { + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + } + // The durable queue cleanup path carries no user identity, so the + // task/session-scoped release must succeed without an owner check. + released, err := repo.DeleteClaimedMessageAttachmentsByTaskSession(ctx, []string{"cleanup-release", "cleanup-release-other"}, "task-1", "session-1") + if err != nil { + t.Fatal(err) + } + if len(released) != 1 || released[0].ID != "cleanup-release" { + t.Fatalf("released = %+v", released) + } + if _, err := repo.GetMessageAttachment(ctx, "cleanup-release"); err == nil { + t.Fatal("released attachment still exists") + } + if _, err := repo.GetMessageAttachment(ctx, "cleanup-release-other"); err != nil { + t.Fatalf("unmatched claim was removed: %v", err) + } + // Replaying the same cleanup is a no-op, not an error. + released, err = repo.DeleteClaimedMessageAttachmentsByTaskSession(ctx, []string{"cleanup-release"}, "task-1", "session-1") + if err != nil { + t.Fatal(err) + } + if len(released) != 0 { + t.Fatalf("replayed release = %+v, want empty", released) + } +} +func TestDeleteClaimedMessageAttachmentsByTaskSessionReleasesTaskScopedRows(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + now := time.Now().UTC() + attachment := &models.TaskMessageAttachment{ + ID: "cleanup-task-scoped", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", + TaskID: "task-1", Name: "task-scoped", MimeType: "text/plain", Kind: "resource", + DeliveryMode: "path", SizeBytes: 1, StorageKey: "cleanup-task-scoped", + State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now, + } + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + + released, err := repo.DeleteClaimedMessageAttachmentsByTaskSession( + ctx, []string{attachment.ID}, "task-1", "session-1", + ) + if err != nil { + t.Fatal(err) + } + if len(released) != 1 || released[0].ID != attachment.ID { + t.Fatalf("released = %+v", released) + } +} + func TestDeleteMessageAttachmentsByTask_RemovesRegistryRows(t *testing.T) { repo := newRepoForSessionTests(t) ctx := context.Background() @@ -340,3 +405,414 @@ func TestDeleteMessageAttachmentsByTask_RemovesRegistryRows(t *testing.T) { t.Fatal("task attachment still exists") } } + +func TestDeleteMessageAttachmentsBySession_RemovesOnlySessionClaims(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + now := time.Now().UTC() + attachments := []*models.TaskMessageAttachment{ + {ID: "session-cleanup", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-session", SessionID: "session-delete", Name: "one", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "session-cleanup", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + {ID: "session-keep", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-session", SessionID: "session-keep", Name: "two", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "session-keep", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + } + for _, attachment := range attachments { + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + } + + removed, err := repo.DeleteMessageAttachmentsBySession(ctx, "task-session", "session-delete") + if err != nil { + t.Fatal(err) + } + if len(removed) != 1 || removed[0].ID != "session-cleanup" { + t.Fatalf("removed = %+v", removed) + } + if _, err := repo.GetMessageAttachment(ctx, "session-cleanup"); err == nil { + t.Fatal("deleted session attachment still exists") + } + if _, err := repo.GetMessageAttachment(ctx, "session-keep"); err != nil { + t.Fatalf("other session attachment was removed: %v", err) + } +} + +func TestTransferMessageAttachments_RebindsClaimedRows(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + now := time.Now().UTC() + attachments := []*models.TaskMessageAttachment{ + {ID: "transfer-one", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-transfer", SessionID: "session-old", Name: "one", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-one", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + {ID: "transfer-other-task", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "other-task", SessionID: "session-old", Name: "two", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-other-task", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + {ID: "transfer-unselected", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-transfer", SessionID: "session-old", Name: "unselected", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-unselected", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + } + for _, attachment := range attachments { + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + } + + if err := repo.TransferMessageAttachments( + ctx, + "task-transfer", + "session-old", + "session-new", + []string{"transfer-one"}, + ); err != nil { + t.Fatal(err) + } + got, err := repo.GetMessageAttachment(ctx, "transfer-one") + if err != nil { + t.Fatal(err) + } + if got.SessionID != "session-new" { + t.Fatalf("transferred attachment session = %q", got.SessionID) + } + other, err := repo.GetMessageAttachment(ctx, "transfer-other-task") + if err != nil { + t.Fatal(err) + } + if other.SessionID != "session-old" { + t.Fatalf("unrelated attachment session = %q", other.SessionID) + } + unselected, err := repo.GetMessageAttachment(ctx, "transfer-unselected") + if err != nil { + t.Fatal(err) + } + if unselected.SessionID != "session-old" { + t.Fatalf("unselected attachment session = %q", unselected.SessionID) + } +} +func TestTransferMessageAttachmentsAcceptsAlreadyTransferredRows(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + now := time.Now().UTC() + for _, attachment := range []*models.TaskMessageAttachment{ + {ID: "transfer-source", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-transfer", SessionID: "session-old", Name: "source", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-source", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + {ID: "transfer-destination", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-transfer", SessionID: "session-new", Name: "destination", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-destination", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + } { + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + } + + if err := repo.TransferMessageAttachments( + ctx, "task-transfer", "session-old", "session-new", + []string{"transfer-source", "transfer-destination"}, + ); err != nil { + t.Fatalf("idempotent attachment transfer: %v", err) + } + for _, attachmentID := range []string{"transfer-source", "transfer-destination"} { + attachment, err := repo.GetMessageAttachment(ctx, attachmentID) + if err != nil { + t.Fatal(err) + } + if attachment.SessionID != "session-new" { + t.Fatalf("attachment %s session = %q, want session-new", attachmentID, attachment.SessionID) + } + } +} + +func TestTransferMessageAttachmentsIgnoresTaskScopedRows(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + now := time.Now().UTC() + attachment := &models.TaskMessageAttachment{ + ID: "transfer-task-scoped", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", + Name: "task-scoped", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", + SizeBytes: 1, StorageKey: "transfer-task-scoped", TaskID: "task-transfer", + State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now, + } + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + + if err := repo.TransferMessageAttachments( + ctx, "task-transfer", "session-old", "session-new", + []string{attachment.ID}, + ); err != nil { + t.Fatalf("task-scoped attachment transfer: %v", err) + } + got, err := repo.GetMessageAttachment(ctx, attachment.ID) + if err != nil { + t.Fatal(err) + } + if got.SessionID != "" { + t.Fatalf("task-scoped attachment session = %q, want empty", got.SessionID) + } +} + +func TestTransferMessageAttachmentsRejectsPartialOwnership(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + now := time.Now().UTC() + for _, attachment := range []*models.TaskMessageAttachment{ + {ID: "transfer-owned", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-transfer", SessionID: "session-old", Name: "owned", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-owned", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + {ID: "transfer-other-session", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", TaskID: "task-transfer", SessionID: "session-other", Name: "other", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-other-session", State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now}, + } { + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + } + + err := repo.TransferMessageAttachments( + ctx, "task-transfer", "session-old", "session-new", + []string{"transfer-owned", "transfer-other-session"}, + ) + + if err == nil { + t.Fatal("partial attachment transfer unexpectedly succeeded") + } + owned, getErr := repo.GetMessageAttachment(ctx, "transfer-owned") + if getErr != nil { + t.Fatal(getErr) + } + if owned.SessionID != "session-old" { + t.Fatalf("partially moved attachment session = %q, want session-old", owned.SessionID) + } +} +func TestTransferMessageAttachmentsAllowsOwnedActiveSessionTransfer(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + if err := repo.CreateTask(ctx, &models.Task{ID: "task-transfer", Title: "Attachment transfer"}); err != nil { + t.Fatal(err) + } + for _, sessionID := range []string{"session-old", "session-new"} { + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: sessionID, TaskID: "task-transfer", State: models.TaskSessionStateWaitingForInput, + }); err != nil { + t.Fatal(err) + } + } + queueRepo, err := messagequeue.NewSQLiteRepository(repo.db, repo.db) + if err != nil { + t.Fatal(err) + } + compensations := queueRepo.(interface { + UpsertSessionTransferCompensation(context.Context, messagequeue.SessionTransferCompensation) error + }) + if err := compensations.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "attachment-transfer", + TaskID: "task-transfer", FromSessionID: "session-old", ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + now := time.Now().UTC() + attachment := &models.TaskMessageAttachment{ + ID: "transfer-owned-active", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", + TaskID: "task-transfer", SessionID: "session-old", Name: "owned", MimeType: "text/plain", + Kind: "resource", DeliveryMode: "path", SizeBytes: 1, StorageKey: "transfer-owned-active", + State: models.AttachmentStateClaimed, ExpiresAt: now.Add(time.Hour), CreatedAt: now, + } + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + + if err := repo.TransferMessageAttachments( + ctx, "task-transfer", "session-old", "session-new", []string{attachment.ID}, + ); err != nil { + t.Fatalf("owned attachment transfer: %v", err) + } + got, err := repo.GetMessageAttachment(ctx, attachment.ID) + if err != nil { + t.Fatal(err) + } + if got.SessionID != "session-new" { + t.Fatalf("transferred attachment session = %q, want session-new", got.SessionID) + } +} +func TestUpsertSessionTransferCompensationValidatesTaskSessions(t *testing.T) { + tests := []struct { + name string + setup func(*Repository, context.Context) (string, string, string) + }{ + { + name: "missing task", + setup: func(_ *Repository, _ context.Context) (string, string, string) { + return "missing-task", "missing-old", "missing-new" + }, + }, + { + name: "missing session", + setup: func(repo *Repository, ctx context.Context) (string, string, string) { + if err := repo.CreateTask(ctx, &models.Task{ID: "task-missing-session", Title: "Transfer"}); err != nil { + t.Fatal(err) + } + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: "session-existing", TaskID: "task-missing-session", + State: models.TaskSessionStateWaitingForInput, + }); err != nil { + t.Fatal(err) + } + return "task-missing-session", "session-existing", "session-missing" + }, + }, + { + name: "session belongs to another task", + setup: func(repo *Repository, ctx context.Context) (string, string, string) { + for _, task := range []*models.Task{ + {ID: "task-owner", Title: "Owner"}, + {ID: "task-other", Title: "Other"}, + } { + if err := repo.CreateTask(ctx, task); err != nil { + t.Fatal(err) + } + } + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: "session-owner", TaskID: "task-owner", + State: models.TaskSessionStateWaitingForInput, + }); err != nil { + t.Fatal(err) + } + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: "session-other", TaskID: "task-other", + State: models.TaskSessionStateWaitingForInput, + }); err != nil { + t.Fatal(err) + } + return "task-owner", "session-owner", "session-other" + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + taskID, fromSessionID, toSessionID := test.setup(repo, ctx) + queueRepo, err := messagequeue.NewSQLiteRepository(repo.db, repo.db) + if err != nil { + t.Fatal(err) + } + persistence := queueRepo.(interface { + UpsertSessionTransferCompensation(context.Context, messagequeue.SessionTransferCompensation) error + }) + err = persistence.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "invalid-transfer-" + test.name, + TaskID: taskID, FromSessionID: fromSessionID, ToSessionID: toSessionID, + }) + if err == nil { + t.Fatal("invalid session transfer compensation unexpectedly succeeded") + } + }) + } +} + +func TestClaimMessageAttachmentsRejectsActiveSessionTransfer(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedWorkspace(t, repo, "workspace-attachments") + if err := repo.CreateTask(ctx, &models.Task{ID: "task-transfer", Title: "Attachment transfer"}); err != nil { + t.Fatal(err) + } + for _, sessionID := range []string{"session-old", "session-new"} { + if err := repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: sessionID, TaskID: "task-transfer", State: models.TaskSessionStateWaitingForInput, + }); err != nil { + t.Fatal(err) + } + } + queueRepo, err := messagequeue.NewSQLiteRepository(repo.db, repo.db) + if err != nil { + t.Fatal(err) + } + compensations := queueRepo.(interface { + UpsertSessionTransferCompensation(context.Context, messagequeue.SessionTransferCompensation) error + }) + if err := compensations.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "claim-transfer", + TaskID: "task-transfer", FromSessionID: "session-old", ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + now := time.Now().UTC() + attachment := &models.TaskMessageAttachment{ + ID: "claim-during-transfer", OwnerID: "owner-1", WorkspaceID: "workspace-attachments", + Name: "during.txt", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", + SizeBytes: 1, StorageKey: "claim-during-transfer", State: models.AttachmentStateStaged, + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + } + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + + err = repo.ClaimMessageAttachments( + ctx, + []string{attachment.ID}, + attachment.OwnerID, + attachment.WorkspaceID, + "task-transfer", + "session-old", + ) + if !errors.Is(err, messagequeue.ErrSessionTransferInProgress) { + t.Fatalf("claim error = %v, want %v", err, messagequeue.ErrSessionTransferInProgress) + } + stored, err := repo.GetMessageAttachment(ctx, attachment.ID) + if err != nil { + t.Fatal(err) + } + if stored.SessionID != "" || stored.State != models.AttachmentStateStaged { + t.Fatalf("attachment after rejected claim = %#v", stored) + } +} + +func TestPostgresAttachmentClaimRejectsActiveSessionTransfer(t *testing.T) { + repoA, repoB, _ := newTaskPostgresRepoPair(t) + ctx := context.Background() + seedWorkspace(t, repoA, "workspace-transfer-claim") + if err := repoA.CreateTask(ctx, &models.Task{ID: "task-transfer", Title: "Attachment transfer"}); err != nil { + t.Fatal(err) + } + for _, sessionID := range []string{"session-old", "session-new"} { + if err := repoA.CreateTaskSession(ctx, &models.TaskSession{ + ID: sessionID, TaskID: "task-transfer", State: models.TaskSessionStateWaitingForInput, + }); err != nil { + t.Fatal(err) + } + } + queueRepo, err := messagequeue.NewSQLiteRepository(repoA.db, repoA.db) + if err != nil { + t.Fatal(err) + } + compensations := queueRepo.(interface { + UpsertSessionTransferCompensation(context.Context, messagequeue.SessionTransferCompensation) error + }) + if err := compensations.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "postgres-claim-transfer", + TaskID: "task-transfer", FromSessionID: "session-old", ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + now := time.Now().UTC() + attachment := &models.TaskMessageAttachment{ + ID: "postgres-claim-during-transfer", OwnerID: "owner-1", WorkspaceID: "workspace-transfer-claim", + Name: "during.txt", MimeType: "text/plain", Kind: "resource", DeliveryMode: "path", + SizeBytes: 1, StorageKey: "postgres-claim-during-transfer", State: models.AttachmentStateStaged, + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + } + if err := repoA.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + + err = repoB.ClaimMessageAttachments( + ctx, + []string{attachment.ID}, + attachment.OwnerID, + attachment.WorkspaceID, + "task-transfer", + "session-old", + ) + if !errors.Is(err, messagequeue.ErrSessionTransferInProgress) { + t.Fatalf("claim error = %v, want %v", err, messagequeue.ErrSessionTransferInProgress) + } + stored, err := repoA.GetMessageAttachment(ctx, attachment.ID) + if err != nil { + t.Fatal(err) + } + if stored.SessionID != "" || stored.State != models.AttachmentStateStaged { + t.Fatalf("attachment after rejected claim = %#v", stored) + } +} diff --git a/apps/backend/internal/task/repository/sqlite/base.go b/apps/backend/internal/task/repository/sqlite/base.go index 6d6f59f875..6188dfdaaa 100644 --- a/apps/backend/internal/task/repository/sqlite/base.go +++ b/apps/backend/internal/task/repository/sqlite/base.go @@ -16,15 +16,16 @@ import ( // Repository provides SQLite-based task storage operations. type Repository struct { - db *sqlx.DB // writer - ro *sqlx.DB // reader (read-only pool) - ownsDB bool - log *logger.Logger - migrate *db.MigrateLogger - queuePurgeMu sync.RWMutex - queuePurger func(context.Context, string) - queuePurgePrepare func(context.Context, string) - queuePurgeNotify func(context.Context, string) + db *sqlx.DB // writer + ro *sqlx.DB // reader (read-only pool) + ownsDB bool + log *logger.Logger + migrate *db.MigrateLogger + queuePurgeMu sync.RWMutex + queuePurger func(context.Context, string) + queuePurgePrepare func(context.Context, string) + queuePurgeNotify func(context.Context, string) + queueSessionPurgeNotify func(context.Context, string, string) // clockNow is a test-only clock seam. Set it before any concurrent // repository call; it carries no synchronization. clockNow func() time.Time @@ -158,6 +159,23 @@ func (r *Repository) notifyTaskQueuePurged(ctx context.Context, taskID string) { } } +// SetTaskSessionQueuePurgeNotifier registers a post-commit observer for queue +// rows removed by DeleteTaskSession. +func (r *Repository) SetTaskSessionQueuePurgeNotifier(notifier func(context.Context, string, string)) { + r.queuePurgeMu.Lock() + defer r.queuePurgeMu.Unlock() + r.queueSessionPurgeNotify = notifier +} + +func (r *Repository) notifyTaskSessionQueuePurged(ctx context.Context, taskID, sessionID string) { + r.queuePurgeMu.RLock() + notifier := r.queueSessionPurgeNotify + r.queuePurgeMu.RUnlock() + if notifier != nil { + notifier(ctx, taskID, sessionID) + } +} + // NewWithDB creates a new SQLite repository with an existing database connection (shared ownership). func NewWithDB(writer, reader *sqlx.DB, log *logger.Logger) (*Repository, error) { return newRepository(writer, reader, log, false) diff --git a/apps/backend/internal/task/repository/sqlite/session.go b/apps/backend/internal/task/repository/sqlite/session.go index 3ca02d673c..92a75e8169 100644 --- a/apps/backend/internal/task/repository/sqlite/session.go +++ b/apps/backend/internal/task/repository/sqlite/session.go @@ -2916,6 +2916,13 @@ func unmarshalSessionSnapshots( return unmarshalSessionJSON(repositorySnapshotJSON, &session.RepositorySnapshot, "repository snapshot") } +// queueSessionLockTablePresent reports whether queue mutations can participate +// in this task transaction. Some repository unit tests intentionally omit the +// message queue schema; production databases always include it. +func (r *Repository) queueSessionLockTablePresent(ctx context.Context) (bool, error) { + return db.TableExistsContext(ctx, r.db, "queue_session_locks") +} + // DeleteTaskSession deletes the exact session incarnation and its pending queue // rows. The session identity fences delayed deletion from removing a replacement // that reuses the same textual session ID. @@ -2939,7 +2946,6 @@ func (r *Repository) DeleteTaskSessionWithAttachments( return nil, err } defer func() { _ = tx.Rollback() }() - if err := r.confirmTaskSessionDeletionIdentityTx(ctx, tx, session); err != nil { return nil, err } @@ -2950,6 +2956,7 @@ func (r *Repository) DeleteTaskSessionWithAttachments( if err := tx.Commit(); err != nil { return nil, err } + r.notifyTaskSessionQueuePurged(ctx, session.TaskID, session.ID) return deletedAttachments, nil } diff --git a/apps/backend/internal/task/repository/sqlite/task_queue_purge_notify_test.go b/apps/backend/internal/task/repository/sqlite/task_queue_purge_notify_test.go index 2a1de7e7e4..7472fc2a7b 100644 --- a/apps/backend/internal/task/repository/sqlite/task_queue_purge_notify_test.go +++ b/apps/backend/internal/task/repository/sqlite/task_queue_purge_notify_test.go @@ -88,6 +88,51 @@ func TestArchiveTaskNotifiesQueuePurgeAfterCommit(t *testing.T) { } } +func TestArchiveTaskPurgesDurableQueueRecovery(t *testing.T) { + repo := newRepoForArchiveTests(t, "task-archive-recovery") + ctx := context.Background() + const sessionID = "session-archive-recovery" + seedLiveSessionForQueue(t, repo, sessionID, "task-archive-recovery") + mqRepo, err := messagequeue.NewSQLiteRepository(repo.db, repo.db) + if err != nil { + t.Fatalf("NewSQLiteRepository: %v", err) + } + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console"}) + if err != nil { + t.Fatalf("logger: %v", err) + } + queue := messagequeue.NewService(mqRepo, messagequeue.DefaultMaxPerSession, log) + entry, err := queue.QueueMessage( + ctx, sessionID, "task-archive-recovery", "follow up", "", "user", false, nil, + ) + if err != nil { + t.Fatal(err) + } + if reserved, ok := queue.ReserveQueued(ctx, sessionID); !ok || reserved.ID != entry.ID { + t.Fatalf("reserved = %#v, ok=%t", reserved, ok) + } + if err := queue.UpsertAttachmentCleanup(ctx, messagequeue.AttachmentCleanup{ + SessionID: sessionID, EntryID: entry.ID, OperationID: "archive-cleanup", + TaskID: entry.TaskID, Attachments: []messagequeue.MessageAttachment{{AttachmentID: "attachment"}}, + }); err != nil { + t.Fatal(err) + } + if err := repo.ArchiveTask(ctx, entry.TaskID); err != nil { + t.Fatal(err) + } + dispatches, err := queue.ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatal(err) + } + cleanups, err := queue.ListAttachmentCleanups(ctx) + if err != nil { + t.Fatal(err) + } + if len(dispatches) != 0 || len(cleanups) != 1 { + t.Fatalf("recovery rows after archive: dispatches=%#v cleanups=%#v", dispatches, cleanups) + } +} + func TestDeleteTaskNotifiesQueuePurgeAfterCommit(t *testing.T) { repo := newRepoForArchiveTests(t, "task-delete-purge-notify") ctx := context.Background() @@ -185,6 +230,161 @@ func TestDeleteTaskSessionPurgesQueuedMessages(t *testing.T) { } } +func TestDeleteTaskSessionPurgesDurableQueueClaimsInTransaction(t *testing.T) { + repo := newRepoForArchiveTests(t, "task-session-recovery-purge") + ctx := context.Background() + const sessionID = "session-recovery-purge" + seedLiveSessionForQueue(t, repo, sessionID, "task-session-recovery-purge") + + mqRepo, err := messagequeue.NewSQLiteRepository(repo.db, repo.db) + if err != nil { + t.Fatalf("NewSQLiteRepository: %v", err) + } + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console"}) + if err != nil { + t.Fatalf("logger: %v", err) + } + queue := messagequeue.NewService(mqRepo, messagequeue.DefaultMaxPerSession, log) + dispatchEntry, err := queue.QueueMessage(ctx, sessionID, "task-session-recovery-purge", "dispatch", "", "user", false, nil) + if err != nil { + t.Fatalf("queue dispatch entry: %v", err) + } + if reserved, ok := queue.ReserveQueued(ctx, sessionID); !ok || reserved.ID != dispatchEntry.ID { + t.Fatalf("reserved dispatch entry = %#v, ok=%t", reserved, ok) + } + sendNowEntry, err := queue.QueueMessage(ctx, sessionID, "task-session-recovery-purge", "send now", "", "user", false, nil) + if err != nil { + t.Fatalf("queue Send Now entry: %v", err) + } + if _, err := mqRepo.ClaimSendNow(ctx, sessionID, []messagequeue.QueuedMessage{*sendNowEntry}); err != nil { + t.Fatalf("claim Send Now entry: %v", err) + } + + if err := deleteTaskSessionByID(ctx, repo, sessionID); err != nil { + t.Fatalf("DeleteTaskSession: %v", err) + } + dispatches, err := mqRepo.(interface { + ListPendingQueueDispatches(context.Context) ([]messagequeue.PendingQueueDispatch, error) + }).ListPendingQueueDispatches(ctx) + if err != nil { + t.Fatalf("list dispatch claims: %v", err) + } + sendNowClaims, err := mqRepo.(interface { + ListPendingSendNowClaims(context.Context) ([]messagequeue.PendingSendNowClaim, error) + }).ListPendingSendNowClaims(ctx) + if err != nil { + t.Fatalf("list Send Now claims: %v", err) + } + if len(dispatches) != 0 || len(sendNowClaims) != 0 { + t.Fatalf("durable claims after session delete: dispatches=%#v send_now=%#v", dispatches, sendNowClaims) + } +} + +func TestDeleteTaskSessionRejectsActiveQueueTransfer(t *testing.T) { + repo := newRepoForArchiveTests(t, "task-session-transfer") + ctx := context.Background() + seedLiveSessionForQueue(t, repo, "session-old", "task-session-transfer") + seedLiveSessionForQueue(t, repo, "session-new", "task-session-transfer") + queueRepo, err := messagequeue.NewSQLiteRepository(repo.db, repo.db) + if err != nil { + t.Fatal(err) + } + compensations := queueRepo.(interface { + UpsertSessionTransferCompensation(context.Context, messagequeue.SessionTransferCompensation) error + }) + if err := compensations.UpsertSessionTransferCompensation(ctx, messagequeue.SessionTransferCompensation{ + OperationID: "session-delete-transfer", + TaskID: "task-session-transfer", FromSessionID: "session-old", ToSessionID: "session-new", + }); err != nil { + t.Fatal(err) + } + + err = deleteTaskSessionByID(ctx, repo, "session-old") + + if !errors.Is(err, messagequeue.ErrSessionTransferInProgress) { + t.Fatalf("DeleteTaskSession error = %v, want %v", err, messagequeue.ErrSessionTransferInProgress) + } + var count int + if err := repo.db.GetContext(ctx, &count, `SELECT COUNT(*) FROM task_sessions WHERE id = 'session-old'`); err != nil { + t.Fatal(err) + } + if count != 1 { + t.Fatalf("session-old row count = %d, want 1", count) + } +} + +func TestDeleteTaskSessionNotifiesQueueSessionPurge(t *testing.T) { + repo := newRepoForArchiveTests(t, "task-session-queue-purge-notify") + ctx := context.Background() + seedLiveSessionForQueue(t, repo, "session-drop", "task-session-queue-purge-notify") + + var notifiedTask, notifiedSession string + repo.SetTaskSessionQueuePurgeNotifier(func(_ context.Context, taskID, sessionID string) { + notifiedTask = taskID + notifiedSession = sessionID + }) + + if err := deleteTaskSessionByID(ctx, repo, "session-drop"); err != nil { + t.Fatalf("DeleteTaskSession: %v", err) + } + if notifiedTask != "task-session-queue-purge-notify" || notifiedSession != "session-drop" { + t.Fatalf("session purge notification = (%q, %q), want (%q, %q)", + notifiedTask, notifiedSession, "task-session-queue-purge-notify", "session-drop") + } +} + +func TestDeleteTaskSessionPurgesPendingMove(t *testing.T) { + repo := newRepoForArchiveTests(t, "task-session-pending-move") + ctx := context.Background() + seedLiveSessionForQueue(t, repo, "session-drop", "task-session-pending-move") + + mqRepo, err := messagequeue.NewSQLiteRepository(repo.db, repo.db) + if err != nil { + t.Fatalf("NewSQLiteRepository: %v", err) + } + if err := mqRepo.SetPendingMove(ctx, "session-drop", &messagequeue.PendingMove{ + TaskID: "task-session-pending-move", + }); err != nil { + t.Fatalf("SetPendingMove: %v", err) + } + + if err := deleteTaskSessionByID(ctx, repo, "session-drop"); err != nil { + t.Fatalf("DeleteTaskSession: %v", err) + } + move, err := mqRepo.GetPendingMove(ctx, "session-drop") + if err != nil { + t.Fatalf("GetPendingMove: %v", err) + } + if move != nil { + t.Fatalf("pending move after session delete = %#v, want nil", move) + } + + err = mqRepo.Insert(ctx, &messagequeue.QueuedMessage{ + SessionID: "session-drop", + TaskID: "task-session-pending-move", + Content: "must reject deleted session", + QueuedBy: "user", + }, 0) + if !errors.Is(err, messagequeue.ErrTaskInactive) { + t.Fatalf("queue admission after session delete error = %v, want ErrTaskInactive", err) + } + if err := mqRepo.SetPendingMove(ctx, "session-drop", &messagequeue.PendingMove{ + TaskID: "task-session-pending-move", + }); !errors.Is(err, messagequeue.ErrTaskInactive) { + t.Fatalf("pending move after session delete error = %v, want ErrTaskInactive", err) + } + err = mqRepo.ReplaceSession(ctx, "session-drop", []messagequeue.QueuedMessage{{ + ID: "restored-after-delete", + SessionID: "session-drop", + TaskID: "task-session-pending-move", + Content: "must not restore", + QueuedBy: "user", + }}, nil) + if !errors.Is(err, messagequeue.ErrTaskInactive) { + t.Fatalf("session restore after delete error = %v, want ErrTaskInactive", err) + } +} + func seedLiveSessionForQueue(t *testing.T, repo *Repository, sessionID, taskID string) { t.Helper() now := time.Now().UTC() @@ -212,3 +412,11 @@ func queueIdentityForSession( SessionIncarnationID: session.QueueIncarnationID, } } + +func deleteTaskSessionByID(ctx context.Context, repo *Repository, sessionID string) error { + session, err := repo.GetTaskSession(ctx, sessionID) + if err != nil { + return err + } + return repo.DeleteTaskSession(ctx, session) +} diff --git a/apps/backend/internal/task/repository/sqlite/task_status_summary_test.go b/apps/backend/internal/task/repository/sqlite/task_status_summary_test.go index 343079ca04..dac9a15b18 100644 --- a/apps/backend/internal/task/repository/sqlite/task_status_summary_test.go +++ b/apps/backend/internal/task/repository/sqlite/task_status_summary_test.go @@ -170,6 +170,12 @@ func TestTaskLastActivityBatch(t *testing.T) { seedTaskAt("task-activity-active", base.Add(2*time.Hour), base.Add(3*time.Hour)) seedTaskAt("task-activity-no-session", base.Add(4*time.Hour), base.Add(5*time.Hour)) seedTaskAt("task-activity-queued", base.Add(4*time.Hour), base.Add(5*time.Hour)) + if _, err := db.Exec(db.Rebind(` + INSERT INTO task_sessions (id, task_id, state, started_at, updated_at) + VALUES (?, ?, 'RUNNING', ?, ?) + `), "session-activity-queued", "task-activity-queued", base, base); err != nil { + t.Fatalf("seed queued activity session: %v", err) + } queueRepo, err := messagequeue.NewSQLiteRepository(db, db) if err != nil { diff --git a/apps/backend/internal/task/service/attachment_service.go b/apps/backend/internal/task/service/attachment_service.go index e3302c2248..839f400ae1 100644 --- a/apps/backend/internal/task/service/attachment_service.go +++ b/apps/backend/internal/task/service/attachment_service.go @@ -293,6 +293,31 @@ func (s *AttachmentService) Release(ctx context.Context, ownerID, taskID, sessio return nil } +type claimedAttachmentCleanupRepository interface { + DeleteClaimedMessageAttachmentsByTaskSession( + context.Context, []string, string, string, + ) ([]*models.TaskMessageAttachment, error) +} + +// ReleaseForCleanup removes claimed descriptors using task/session ownership +// rather than a user identity. It is restricted to durable queue cleanup. +func (s *AttachmentService) ReleaseForCleanup( + ctx context.Context, taskID, sessionID string, ids []string, +) error { + repo, ok := s.repo.(claimedAttachmentCleanupRepository) + if !ok { + return errors.New("attachment cleanup release is unavailable") + } + attachments, err := repo.DeleteClaimedMessageAttachmentsByTaskSession(ctx, ids, taskID, sessionID) + if err != nil { + return err + } + for _, attachment := range attachments { + s.removeBytes(attachment) + } + return nil +} + // DeleteByTask removes all attachment registry rows and private bytes owned by // a task. Task deletion must clean claimed rows as well as staged rows because // only staged rows participate in expiry maintenance. @@ -307,6 +332,31 @@ func (s *AttachmentService) DeleteByTask(ctx context.Context, taskID string) err return nil } +// DeleteBySession removes all claimed attachment descriptors and private bytes +// owned by a deleted task session. Staged uploads are intentionally excluded +// because they are not bound to a session and expire independently. +func (s *AttachmentService) DeleteBySession(ctx context.Context, taskID, sessionID string) error { + attachments, err := s.repo.DeleteMessageAttachmentsBySession(ctx, taskID, sessionID) + if err != nil { + return err + } + for _, attachment := range attachments { + s.removeBytes(attachment) + } + return nil +} + +// TransferSession rebinds only the claimed prompt attachments represented by +// the queue transfer operation. The source-session predicate is the CAS guard +// used by both forward transfer and rollback. +func (s *AttachmentService) TransferSession( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + attachmentIDs []string, +) error { + return s.repo.TransferMessageAttachments(ctx, taskID, oldSessionID, newSessionID, attachmentIDs) +} + type transactionalWorkspaceAttachmentRepository interface { DeleteMessageAttachmentsByWorkspaceTx( context.Context, diff --git a/apps/backend/internal/task/service/service_attachments.go b/apps/backend/internal/task/service/service_attachments.go index 3a14fcb6fc..311d38d3f7 100644 --- a/apps/backend/internal/task/service/service_attachments.go +++ b/apps/backend/internal/task/service/service_attachments.go @@ -98,3 +98,106 @@ func (s *Service) ReleaseMessageAttachments(ctx context.Context, taskID, session } return s.attachmentSvc.Release(ctx, identity.UserID, taskID, sessionID, ids) } + +// ResolveMessageAttachmentOwner recovers the authenticated owner needed to +// replay a legacy durable queue cleanup after its owner column was introduced. +func (s *Service) ResolveMessageAttachmentOwner( + ctx context.Context, taskID, sessionID string, attachments []v1.MessageAttachment, +) (string, error) { + if s.attachmentSvc == nil { + return "", errors.New("file-backed attachments are unavailable") + } + ids := make([]string, 0, len(attachments)) + seen := make(map[string]struct{}, len(attachments)) + for _, attachment := range attachments { + if attachment.AttachmentID == "" { + continue + } + if _, ok := seen[attachment.AttachmentID]; ok { + continue + } + seen[attachment.AttachmentID] = struct{}{} + ids = append(ids, attachment.AttachmentID) + } + if len(ids) == 0 { + return "", nil + } + rows, err := s.attachmentSvc.repo.ListMessageAttachments(ctx, ids) + if err != nil { + return "", err + } + if len(rows) != len(ids) { + return "", errors.New("durable cleanup attachment is missing") + } + ownerID := "" + for _, attachment := range rows { + attachmentOwner, err := validateMessageAttachmentOwner( + attachment, taskID, sessionID, + ) + if err != nil { + return "", err + } + if ownerID == "" { + ownerID = attachmentOwner + } else if ownerID != attachmentOwner { + return "", errors.New("durable cleanup attachments have different owners") + } + } + return ownerID, nil +} + +func validateMessageAttachmentOwner( + attachment *models.TaskMessageAttachment, taskID, sessionID string, +) (string, error) { + // Legacy claim-pending cleanups may still reference staged descriptors. + if attachment == nil || attachment.TaskID != taskID || + (attachment.State != models.AttachmentStateClaimed && + attachment.State != models.AttachmentStateStaged) || + (attachment.SessionID != "" && attachment.SessionID != sessionID) || + attachment.OwnerID == "" { + return "", errors.New("durable cleanup attachment ownership is invalid") + } + return attachment.OwnerID, nil +} + +// ReleaseMessageAttachmentsForCleanup provides the durable queue retry path +// with a task/session-scoped release that does not require user auth. +func (s *Service) ReleaseMessageAttachmentsForCleanup( + ctx context.Context, taskID, sessionID string, attachments []v1.MessageAttachment, +) error { + if s.attachmentSvc == nil { + return errors.New("file-backed attachments are unavailable") + } + ids := make([]string, 0, len(attachments)) + for _, attachment := range attachments { + if attachment.AttachmentID != "" { + ids = append(ids, attachment.AttachmentID) + } + } + return s.attachmentSvc.ReleaseForCleanup(ctx, taskID, sessionID, ids) +} + +// DeleteSessionMessageAttachments removes file-backed prompt attachments +// claimed by a task session that has already been deleted. +func (s *Service) DeleteSessionMessageAttachments(ctx context.Context, taskID, sessionID string) error { + if s.attachmentSvc == nil { + return nil + } + return s.attachmentSvc.DeleteBySession(ctx, taskID, sessionID) +} + +// TransferSessionMessageAttachments keeps claimed prompt attachments aligned +// with a queued session transfer without touching unrelated destination claims. +func (s *Service) TransferSessionMessageAttachments( + ctx context.Context, + taskID, oldSessionID, newSessionID string, + attachmentIDs []string, +) error { + if s.attachmentSvc == nil { + if len(attachmentIDs) == 0 { + return nil + } + return errors.New("file-backed attachments are unavailable") + } + return s.attachmentSvc.TransferSession(ctx, taskID, oldSessionID, newSessionID, attachmentIDs) +} diff --git a/apps/backend/internal/task/service/service_attachments_cleanup_test.go b/apps/backend/internal/task/service/service_attachments_cleanup_test.go new file mode 100644 index 0000000000..c4f56e2333 --- /dev/null +++ b/apps/backend/internal/task/service/service_attachments_cleanup_test.go @@ -0,0 +1,91 @@ +package service + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/kandev/kandev/internal/task/models" + v1 "github.com/kandev/kandev/pkg/api/v1" +) + +// The durable queue cleanup path carries no user identity: ReleaseForCleanup +// must remove claimed rows and bytes scoped by task/session alone. +func TestReleaseForCleanupRemovesClaimedRowsAndBytes(t *testing.T) { + svc, repo, root, _ := newAttachmentTestService(t) + ctx := context.Background() + released := stageTestAttachment(t, svc, "user-a", "released.png", "x") + kept := stageTestAttachment(t, svc, "user-a", "kept.png", "y") + if err := svc.Claim(ctx, "user-a", "ws-att", "task-1", "sess-1", []string{released.ID, kept.ID}); err != nil { + t.Fatalf("Claim: %v", err) + } + if err := svc.ReleaseForCleanup(ctx, "task-1", "sess-1", []string{released.ID}); err != nil { + t.Fatalf("ReleaseForCleanup: %v", err) + } + if _, err := repo.GetMessageAttachment(ctx, released.ID); !errors.Is(err, ErrAttachmentNotFound) { + t.Fatalf("released row = %v, want ErrAttachmentNotFound", err) + } + if _, err := os.Stat(filepath.Join(root, "attachments", released.StorageKey)); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("released bytes still present: %v", err) + } + if _, err := repo.GetMessageAttachment(ctx, kept.ID); err != nil { + t.Fatalf("unreferenced release must not touch other rows: %v", err) + } +} + +// ResolveMessageAttachmentOwner recovers the claim owner so a legacy durable +// cleanup (persisted before the owner column) can replay under its identity. +func TestResolveMessageAttachmentOwnerRecoversClaimOwner(t *testing.T) { + attachmentSvc, _, _, _ := newAttachmentTestService(t) + svc := &Service{attachmentSvc: attachmentSvc} + ctx := context.Background() + first := stageTestAttachment(t, attachmentSvc, "user-a", "first.png", "x") + second := stageTestAttachment(t, attachmentSvc, "user-a", "second.png", "y") + if err := attachmentSvc.Claim(ctx, "user-a", "ws-att", "task-1", "sess-1", []string{first.ID, second.ID}); err != nil { + t.Fatalf("Claim: %v", err) + } + owner, err := svc.ResolveMessageAttachmentOwner(ctx, "task-1", "sess-1", []v1.MessageAttachment{ + {AttachmentID: first.ID}, {AttachmentID: second.ID}, + }) + if err != nil { + t.Fatalf("ResolveMessageAttachmentOwner: %v", err) + } + if owner != "user-a" { + t.Fatalf("owner = %q, want user-a", owner) + } + if _, err := svc.ResolveMessageAttachmentOwner(ctx, "task-1", "sess-other", []v1.MessageAttachment{ + {AttachmentID: first.ID}, + }); err == nil || !strings.Contains(err.Error(), "ownership is invalid") { + t.Fatalf("foreign-session resolve = %v, want ownership error", err) + } +} +func TestResolveMessageAttachmentOwnerRecoversStagedClaimOwner(t *testing.T) { + attachmentSvc, repo, _, _ := newAttachmentTestService(t) + svc := &Service{attachmentSvc: attachmentSvc} + ctx := context.Background() + attachment := &models.TaskMessageAttachment{ + ID: "staged-owner-recovery", OwnerID: "user-staged", WorkspaceID: "ws-att", + TaskID: "task-1", SessionID: "sess-1", Name: "staged.png", + MimeType: "image/png", Kind: "image", DeliveryMode: "path", + SizeBytes: 1, StorageKey: "staged-owner-recovery", + State: models.AttachmentStateStaged, ExpiresAt: time.Now().Add(time.Hour), + CreatedAt: time.Now(), + } + if err := repo.CreateMessageAttachment(ctx, attachment); err != nil { + t.Fatal(err) + } + + owner, err := svc.ResolveMessageAttachmentOwner(ctx, "task-1", "sess-1", []v1.MessageAttachment{ + {AttachmentID: attachment.ID}, + }) + if err != nil { + t.Fatalf("ResolveMessageAttachmentOwner: %v", err) + } + if owner != "user-staged" { + t.Fatalf("owner = %q, want user-staged", owner) + } +} diff --git a/apps/backend/internal/task/service/service_attachments_transfer_test.go b/apps/backend/internal/task/service/service_attachments_transfer_test.go new file mode 100644 index 0000000000..9b13eda3b5 --- /dev/null +++ b/apps/backend/internal/task/service/service_attachments_transfer_test.go @@ -0,0 +1,17 @@ +package service + +import ( + "context" + "testing" +) + +func TestTransferSessionMessageAttachmentsRejectsMissingAttachmentService(t *testing.T) { + service := &Service{} + err := service.TransferSessionMessageAttachments( + context.Background(), "task", "session-old", "session-new", + []string{"attachment"}, + ) + if err == nil { + t.Fatal("transfer unexpectedly succeeded without attachment service") + } +} diff --git a/apps/backend/pkg/websocket/actions.go b/apps/backend/pkg/websocket/actions.go index c8b5f5a84a..ebcc01d5e4 100644 --- a/apps/backend/pkg/websocket/actions.go +++ b/apps/backend/pkg/websocket/actions.go @@ -142,6 +142,9 @@ const ( ActionMessageQueueCancel = "message.queue.cancel" // Clears the entire queue for a session ActionMessageQueueGet = "message.queue.get" ActionMessageQueueUpdate = "message.queue.update" + ActionMessageQueueEditBegin = "message.queue.edit.begin" + ActionMessageQueueEditRenew = "message.queue.edit.renew" + ActionMessageQueueEditEnd = "message.queue.edit.end" ActionMessageQueueAppend = "message.queue.append" ActionMessageQueueDrain = "message.queue.drain" // Dispatch one queued entry now when the session is promptable ActionMessageQueueSendNow = "message.queue.send_now" // Interrupt and replace the active turn with an exact queue selection diff --git a/apps/backend/pkg/websocket/message.go b/apps/backend/pkg/websocket/message.go index 6947f62941..ac350cd5d7 100644 --- a/apps/backend/pkg/websocket/message.go +++ b/apps/backend/pkg/websocket/message.go @@ -2,10 +2,25 @@ package websocket import ( + "context" "encoding/json" "time" ) +type connectionIDContextKey struct{} + +// WithConnectionID binds an inbound message to its live WebSocket connection. +// The value is server supplied and cannot be forged through a message payload. +func WithConnectionID(ctx context.Context, id string) context.Context { + return context.WithValue(ctx, connectionIDContextKey{}, id) +} + +// ConnectionID returns the server-assigned WebSocket connection identifier. +func ConnectionID(ctx context.Context) string { + id, _ := ctx.Value(connectionIDContextKey{}).(string) + return id +} + // MessageType represents the type of WebSocket message type MessageType string diff --git a/apps/web/components/task/chat/queued-ghost-edit-view.tsx b/apps/web/components/task/chat/queued-ghost-edit-view.tsx new file mode 100644 index 0000000000..0645679782 --- /dev/null +++ b/apps/web/components/task/chat/queued-ghost-edit-view.tsx @@ -0,0 +1,85 @@ +"use client"; + +import type { KeyboardEvent, RefObject } from "react"; +import { IconCheck } from "@tabler/icons-react"; +import { useTranslation } from "react-i18next"; +import { Button } from "@kandev/ui"; +import { Textarea } from "@kandev/ui/textarea"; +import { cn } from "@/lib/utils"; +import { AttachmentRow, type QueuedAttachment } from "@/components/task/chat/queued-attachment-row"; + +type QueuedGhostEditViewProps = { + value: string; + saving: boolean; + attachments: QueuedAttachment[]; + onChange: (value: string) => void; + onSave: () => void; + onCancel: () => void; + textareaRef: RefObject; +}; + +export function QueuedGhostEditView({ + value, + saving, + attachments, + onChange, + onSave, + onCancel, + textareaRef, +}: QueuedGhostEditViewProps) { + const { t } = useTranslation(); + const onKeyDown = (event: KeyboardEvent) => { + if (event.key === "Escape") { + event.preventDefault(); + // Claim the key here: once this edit is cancelled, nothing further up + // the tree (e.g. a clarification panel's own Escape-collapses handler) + // should also react to the same keypress. + event.stopPropagation(); + onCancel(); + } else if (event.key === "Enter" && (event.metaKey || event.ctrlKey)) { + event.preventDefault(); + onSave(); + } + }; + return ( +
+ +