3751 lines
131 KiB
Go
3751 lines
131 KiB
Go
// Package messagequeue persists per-session FIFO message queues and deferred
|
|
// workflow moves. The queue lets a user (or another agent via the
|
|
// message_task_kandev MCP tool) enqueue follow-up prompts while the target
|
|
// session is still busy with a turn; handleAgentReady drains one entry per
|
|
// turn after each turn completes.
|
|
//
|
|
// Storage is abstracted behind Repository (SQLite in production, in-memory
|
|
// for tests). The service enforces a per-session cap (default 10) and surfaces
|
|
// overflow as ErrQueueFull rather than silently dropping the new message.
|
|
package messagequeue
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"reflect"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/kandev/kandev/internal/common/logger"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type autoMergePolicyLoader struct {
|
|
load func(context.Context) (bool, int64, error)
|
|
}
|
|
|
|
// Service manages queued messages for sessions, backed by Repository.
|
|
type Service struct {
|
|
repo Repository
|
|
maxPerSession atomic.Int64
|
|
mergeEnabled atomic.Bool
|
|
autoMergePolicy atomic.Pointer[AutoMergePolicy]
|
|
autoMergeLoader atomic.Pointer[autoMergePolicyLoader]
|
|
statusEpoch string
|
|
logger *logger.Logger
|
|
// 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 {
|
|
mu sync.Mutex
|
|
refs int
|
|
}
|
|
|
|
type sessionAdmissionContextKey struct{}
|
|
|
|
type sessionAdmissionToken struct {
|
|
service *Service
|
|
sessionID string
|
|
afterLifecycleUnlock *[]func()
|
|
}
|
|
|
|
type queueAttachmentRepository interface {
|
|
InsertForSessionWithClaim(context.Context, QueueSessionIdentity, *QueuedMessage, QueueAttachmentClaim, int) error
|
|
AutoMergeCandidateIntoAboveForSessionWithClaim(context.Context, QueueSessionIdentity, *QueuedMessage, QueueAttachmentClaim) (*QueuedMessage, bool, error)
|
|
UpdateContentAndMetadataForSessionWithClaim(context.Context, QueueSessionIdentity, string, string, []MessageAttachment, map[string]interface{}, string, QueueAttachmentClaim) error
|
|
}
|
|
|
|
type deletedSessionRepository interface {
|
|
PurgeDeletedSession(context.Context, QueueSessionIdentity) error
|
|
}
|
|
|
|
type autoMergeAdmissionRepository interface {
|
|
InsertForSessionWithPolicy(
|
|
context.Context,
|
|
QueueSessionIdentity,
|
|
*QueuedMessage,
|
|
*QueueAttachmentClaim,
|
|
int,
|
|
AutoMergePolicy,
|
|
) error
|
|
AutoMergeCandidateIntoAboveForSessionWithPolicy(
|
|
context.Context,
|
|
QueueSessionIdentity,
|
|
*QueuedMessage,
|
|
*QueueAttachmentClaim,
|
|
AutoMergePolicy,
|
|
) (*QueuedMessage, bool, error)
|
|
AutoMergeIntoAboveForSessionWithPolicy(
|
|
context.Context,
|
|
QueueSessionIdentity,
|
|
string,
|
|
AutoMergePolicy,
|
|
) (*QueuedMessage, bool, error)
|
|
}
|
|
|
|
// NewService creates a Service backed by the supplied repository. maxPerSession
|
|
// is the per-session cap (entries beyond this return ErrQueueFull on insert);
|
|
// 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),
|
|
editLeases: make(map[editLeaseKey]*QueueEditLease),
|
|
editRevisions: make(map[editLeaseKey]int64),
|
|
editRevisionTaskIDs: make(map[editLeaseKey]string),
|
|
}
|
|
service.SetMaxPerSession(maxPerSession)
|
|
service.mergeEnabled.Store(true)
|
|
service.SetAutoMergePolicy(true, 0)
|
|
return service
|
|
}
|
|
|
|
// NewServiceMemory returns a Service backed by an in-memory repository.
|
|
// Convenience constructor for tests; cap defaults to 10 for parity with prod.
|
|
func NewServiceMemory(log *logger.Logger) *Service {
|
|
return NewService(NewMemoryRepository(), DefaultMaxPerSession, log)
|
|
}
|
|
|
|
// ResolveSessionIdentity returns the repository-authoritative immutable queue
|
|
// identity for trusted internal callers migrating legacy session snapshots.
|
|
func (s *Service) ResolveSessionIdentity(ctx context.Context, taskID, sessionID string) (QueueSessionIdentity, error) {
|
|
return s.repo.ResolveSessionIdentity(ctx, taskID, sessionID)
|
|
}
|
|
|
|
// SupportsAtomicDeferredMoveTransition reports whether task and queue rows share
|
|
// one SQL transaction. The in-memory queue cannot participate in task storage.
|
|
func (s *Service) SupportsAtomicDeferredMoveTransition() bool {
|
|
_, ok := s.repo.(*sqliteRepository)
|
|
return ok
|
|
}
|
|
|
|
// MaxPerSession returns the configured per-session cap.
|
|
func (s *Service) MaxPerSession() int { return int(s.maxPerSession.Load()) }
|
|
|
|
// LifecycleGeneration returns the task archive/delete generation captured by
|
|
// send-now restoration guards.
|
|
func (s *Service) LifecycleGeneration(ctx context.Context, taskID string) (int64, error) {
|
|
return s.repo.LifecycleGeneration(ctx, taskID)
|
|
}
|
|
|
|
// ListDurableLifecycleEntries returns lifecycle rows that remain persisted
|
|
// until executor acceptance. Startup recovery uses it to reconcile queue rows
|
|
// with their owning automation reservations.
|
|
func (s *Service) ListDurableLifecycleEntries(ctx context.Context) ([]QueuedMessage, error) {
|
|
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)
|
|
}
|
|
|
|
// ListDurableDeliveryEntries returns all durable queue receipts used by
|
|
// startup reconciliation.
|
|
func (s *Service) ListDurableDeliveryEntries(ctx context.Context) ([]QueuedMessage, error) {
|
|
return s.repo.ListDurableDeliveryEntries(ctx)
|
|
}
|
|
|
|
// SetMaxPerSession applies a new admission cap without pruning existing rows.
|
|
// Non-positive values disable the cap.
|
|
func (s *Service) SetMaxPerSession(maxPerSession int) {
|
|
if maxPerSession < 0 {
|
|
maxPerSession = 0
|
|
}
|
|
s.maxPerSession.Store(int64(maxPerSession))
|
|
}
|
|
|
|
// MergeEnabled reports whether MergeIntoAbove is currently allowed. Enabled
|
|
// by default; see SetMergeEnabled.
|
|
func (s *Service) MergeEnabled() bool { return s.mergeEnabled.Load() }
|
|
|
|
// SetMergeEnabled toggles whether queued messages may be folded into the
|
|
// entry above them via MergeIntoAbove. Disabling it does not affect merges
|
|
// already applied.
|
|
func (s *Service) SetMergeEnabled(enabled bool) {
|
|
s.mergeEnabled.Store(enabled)
|
|
}
|
|
|
|
// AutoMergeEnabled reports the global admission policy used by sessions that
|
|
// do not have an explicit override.
|
|
func (s *Service) AutoMergeEnabled() bool {
|
|
return s.globalAutoMergePolicy().Enabled
|
|
}
|
|
|
|
// SetAutoMergeEnabled updates the global policy while preserving the legacy
|
|
// boolean target used by callers that do not own the persisted revision.
|
|
func (s *Service) SetAutoMergeEnabled(enabled bool) {
|
|
current := s.globalAutoMergePolicy()
|
|
revision := current.Revision
|
|
if current.Enabled != enabled {
|
|
revision++
|
|
}
|
|
s.SetAutoMergePolicy(enabled, revision)
|
|
}
|
|
|
|
// SetAutoMergePolicy atomically publishes a complete global policy snapshot.
|
|
func (s *Service) SetAutoMergePolicy(enabled bool, revision int64) {
|
|
policy := &AutoMergePolicy{
|
|
Enabled: enabled,
|
|
Source: AutoMergeSourceGlobal,
|
|
Revision: revision,
|
|
}
|
|
s.autoMergePolicy.Store(policy)
|
|
}
|
|
|
|
// SetAutoMergePolicyLoader installs the shared persisted-policy reader used by
|
|
// identity-bound admissions and snapshots. The atomic cache remains available
|
|
// to legacy callers that cannot carry a context.
|
|
func (s *Service) SetAutoMergePolicyLoader(
|
|
load func(context.Context) (bool, int64, error),
|
|
) {
|
|
if load == nil {
|
|
s.autoMergeLoader.Store(nil)
|
|
return
|
|
}
|
|
s.autoMergeLoader.Store(&autoMergePolicyLoader{load: load})
|
|
}
|
|
|
|
func (s *Service) globalAutoMergePolicy() AutoMergePolicy {
|
|
policy := s.autoMergePolicy.Load()
|
|
if policy == nil {
|
|
return AutoMergePolicy{Enabled: true, Source: AutoMergeSourceGlobal}
|
|
}
|
|
return *policy
|
|
}
|
|
|
|
// ResolveAutoMergePolicy returns the session override when present, otherwise
|
|
// the current global policy snapshot.
|
|
func (s *Service) ResolveAutoMergePolicy(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
) (AutoMergePolicy, error) {
|
|
override, err := s.repo.GetAutoMergeOverride(ctx, identity)
|
|
if err != nil {
|
|
return AutoMergePolicy{}, err
|
|
}
|
|
if override == nil {
|
|
return s.loadGlobalAutoMergePolicy(ctx)
|
|
}
|
|
return AutoMergePolicy{
|
|
Enabled: override.Enabled,
|
|
Source: AutoMergeSourceSession,
|
|
Revision: override.Revision,
|
|
}, nil
|
|
}
|
|
|
|
func (s *Service) loadGlobalAutoMergePolicy(ctx context.Context) (AutoMergePolicy, error) {
|
|
if loader := s.autoMergeLoader.Load(); loader != nil {
|
|
enabled, revision, err := loader.load(ctx)
|
|
if err != nil {
|
|
return AutoMergePolicy{}, err
|
|
}
|
|
s.SetAutoMergePolicy(enabled, revision)
|
|
}
|
|
return s.globalAutoMergePolicy(), nil
|
|
}
|
|
|
|
// SetSessionAutoMerge persists and returns an explicit per-session policy.
|
|
func (s *Service) SetSessionAutoMerge(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
enabled bool,
|
|
) (AutoMergePolicy, error) {
|
|
var override AutoMergeOverride
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var setErr error
|
|
override, setErr = s.repo.SetAutoMergeOverride(admittedCtx, identity, enabled)
|
|
return setErr
|
|
})
|
|
if err != nil {
|
|
return AutoMergePolicy{}, err
|
|
}
|
|
return AutoMergePolicy{
|
|
Enabled: override.Enabled,
|
|
Source: AutoMergeSourceSession,
|
|
Revision: override.Revision,
|
|
}, nil
|
|
}
|
|
|
|
// WithSessionAdmission runs fn under the per-session queue admission lock.
|
|
// All queue insertion and mutation paths use the same lock. The callback must
|
|
// complete synchronously; queue methods called with its context reuse the held
|
|
// lock.
|
|
func (s *Service) WithSessionAdmission(ctx context.Context, sessionID string, fn func(context.Context) error) error {
|
|
if fn == nil {
|
|
return errors.New("session admission callback is nil")
|
|
}
|
|
if token, ok := ctx.Value(sessionAdmissionContextKey{}).(sessionAdmissionToken); ok &&
|
|
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 {
|
|
entry = &sessionAdmission{}
|
|
s.admissions[sessionID] = entry
|
|
}
|
|
entry.refs++
|
|
s.admissionMu.Unlock()
|
|
|
|
entry.mu.Lock()
|
|
defer func() {
|
|
entry.mu.Unlock()
|
|
s.admissionMu.Lock()
|
|
entry.refs--
|
|
if entry.refs == 0 && s.admissions[sessionID] == entry {
|
|
delete(s.admissions, sessionID)
|
|
}
|
|
s.admissionMu.Unlock()
|
|
}()
|
|
|
|
admittedCtx := context.WithValue(ctx, sessionAdmissionContextKey{}, sessionAdmissionToken{
|
|
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) {
|
|
return s.QueueMessageWithMetadata(ctx, sessionID, taskID, content, model, userID, planMode, attachments, nil)
|
|
}
|
|
|
|
// QueueMessageWithMetadata is like QueueMessage but stores extra metadata that
|
|
// is propagated to the resulting Message row when the queued message is
|
|
// drained (e.g. sender_task_id for messages sent via message_task_kandev).
|
|
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(
|
|
ctx, &identity, sessionID, taskID, content, model, userID, planMode, attachments, metadata, nil, nil,
|
|
)
|
|
}
|
|
|
|
// QueueMessageWithMetadataAfterInsert admits an exact source row, runs
|
|
// afterInsert while the per-session admission lock remains held, then attempts
|
|
// the snapshotted automatic fold. Callers use the hook to claim staged
|
|
// attachments before folding; if it returns an error, no fold is attempted.
|
|
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(
|
|
ctx, &identity, sessionID, taskID, content, model, userID, planMode, attachments, metadata, nil, afterInsert,
|
|
)
|
|
}
|
|
|
|
// QueueMessageWithMetadataForSession admits a browser message using the
|
|
// effective policy bound to the supplied immutable session identity.
|
|
func (s *Service) QueueMessageWithMetadataForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
) (*QueuedMessage, error) {
|
|
return s.queueMessageWithMetadataAdmission(
|
|
ctx, &identity, identity.SessionID, identity.TaskID, content, model, userID, planMode, attachments, metadata, nil, nil,
|
|
)
|
|
}
|
|
|
|
// QueueMessageWithMetadataForSessionWithClientQueueID admits an ordinary
|
|
// browser queue message with a durable caller-owned replay identity. The
|
|
// receipt and any queue insertion, fold, or attachment claim share one
|
|
// repository transaction.
|
|
func (s *Service) QueueMessageWithMetadataForSessionWithClientQueueID(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
clientQueueID, content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
claim *QueueAttachmentClaim,
|
|
) (*QueuedMessage, bool, error) {
|
|
if clientQueueID == "" || len(clientQueueID) > MaxQueueAdmissionIDLength {
|
|
return nil, false, errors.New("client queue id is invalid")
|
|
}
|
|
writer, ok := s.repo.(queueAdmissionRepository)
|
|
if !ok {
|
|
return nil, false, ErrQueueAdmissionUnavailable
|
|
}
|
|
candidate := &QueuedMessage{
|
|
ID: clientQueueID, SessionID: identity.SessionID, TaskID: identity.TaskID,
|
|
Content: content, Model: model, PlanMode: planMode,
|
|
Attachments: append([]MessageAttachment(nil), attachments...),
|
|
Metadata: copyMessageMetadata(metadata, 0), QueuedBy: userID,
|
|
}
|
|
var admitted *QueuedMessage
|
|
var replay bool
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
for {
|
|
policy := s.resolveAdmissionAutoMergePolicy(admittedCtx, &identity, identity.SessionID)
|
|
var err error
|
|
admitted, replay, err = writer.AdmitQueueMessage(
|
|
admittedCtx, identity, clientQueueID, candidate, claim, s.MaxPerSession(), policy,
|
|
)
|
|
if errors.Is(err, ErrAutoMergePolicyChanged) {
|
|
continue
|
|
}
|
|
return err
|
|
}
|
|
})
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
return admitted, replay, nil
|
|
}
|
|
|
|
// LookupQueueAdmissionWithClientQueueID checks a caller-owned admission
|
|
// receipt before validation that depends on mutable external references or
|
|
// attachment ownership.
|
|
func (s *Service) LookupQueueAdmissionWithClientQueueID(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
clientQueueID, content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
) (*QueuedMessage, bool, error) {
|
|
if clientQueueID == "" || len(clientQueueID) > MaxQueueAdmissionIDLength {
|
|
return nil, false, errors.New("client queue id is invalid")
|
|
}
|
|
reader, ok := s.repo.(queueAdmissionRepository)
|
|
if !ok {
|
|
return nil, false, ErrQueueAdmissionUnavailable
|
|
}
|
|
candidate := &QueuedMessage{
|
|
ID: clientQueueID, SessionID: identity.SessionID, TaskID: identity.TaskID,
|
|
Content: content, Model: model, PlanMode: planMode,
|
|
Attachments: append([]MessageAttachment(nil), attachments...),
|
|
Metadata: copyMessageMetadata(metadata, 0), QueuedBy: userID,
|
|
}
|
|
var admitted *QueuedMessage
|
|
var replay bool
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var lookupErr error
|
|
admitted, replay, lookupErr = reader.LookupQueueAdmission(admittedCtx, identity, clientQueueID, candidate)
|
|
return lookupErr
|
|
})
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
return admitted, replay, nil
|
|
}
|
|
|
|
// QueueMessageWithMetadataForSessionWithClaim atomically claims staged
|
|
// attachments when admission inserts a queue entry. An empty claim carries
|
|
// inline-only attachments and may use the compatible full-queue fold path.
|
|
func (s *Service) QueueMessageWithMetadataForSessionWithClaim(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
claim QueueAttachmentClaim,
|
|
) (*QueuedMessage, error) {
|
|
return s.queueMessageWithMetadataAdmission(
|
|
ctx, &identity, identity.SessionID, identity.TaskID, content, model, userID, planMode, attachments, metadata, &claim, nil,
|
|
)
|
|
}
|
|
|
|
// QueueMessageWithMetadataForSessionAfterInsert is the identity-bound variant
|
|
// used while attachment claiming still participates in the admission lock.
|
|
func (s *Service) QueueMessageWithMetadataForSessionAfterInsert(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
afterInsert func(context.Context, *QueuedMessage) error,
|
|
) (*QueuedMessage, error) {
|
|
return s.queueMessageWithMetadataAdmission(
|
|
ctx, &identity, identity.SessionID, identity.TaskID, content, model, userID, planMode, attachments, metadata, nil, afterInsert,
|
|
)
|
|
}
|
|
|
|
// queueMessageWithMetadataAdmission resolves policy and completes admission
|
|
// under one per-session lock. Identity-bound repositories compare that policy
|
|
// inside every durable mutation; an override race restarts before admission.
|
|
//
|
|
//nolint:gocognit // retry boundaries keep policy validation before each mutation.
|
|
func (s *Service) queueMessageWithMetadataAdmission(ctx context.Context, identity *QueueSessionIdentity, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, claim *QueueAttachmentClaim, afterInsert func(context.Context, *QueuedMessage) error) (*QueuedMessage, error) {
|
|
var queued *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
for {
|
|
policy := s.resolveAdmissionAutoMergePolicy(admittedCtx, identity, sessionID)
|
|
source, insertErr := s.insertQueueMessageWithMetadata(
|
|
admittedCtx, identity, sessionID, taskID, content, model, userID, planMode,
|
|
attachments, metadata, claim, s.MaxPerSession(), policy,
|
|
)
|
|
if errors.Is(insertErr, ErrAutoMergePolicyChanged) {
|
|
continue
|
|
}
|
|
autoMergeEnabled := policy != nil && policy.Enabled
|
|
if insertErr != nil {
|
|
hasStagedAttachments := claim != nil && len(claim.IDs) > 0
|
|
if !errors.Is(insertErr, ErrQueueFull) || !autoMergeEnabled || hasStagedAttachments || afterInsert != nil {
|
|
return insertErr
|
|
}
|
|
var merged *QueuedMessage
|
|
source, merged, insertErr = s.admitQueueFullMessage(
|
|
admittedCtx, identity, insertErr, sessionID, taskID, content, model, userID, planMode,
|
|
attachments, metadata, claim, s.MaxPerSession(), policy,
|
|
)
|
|
if errors.Is(insertErr, ErrAutoMergePolicyChanged) {
|
|
continue
|
|
}
|
|
if insertErr != nil {
|
|
return insertErr
|
|
}
|
|
if merged != nil {
|
|
queued = merged
|
|
return nil
|
|
}
|
|
}
|
|
if afterInsert != nil {
|
|
if err := afterInsert(admittedCtx, source); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
queued = s.finalizeAutoMerge(admittedCtx, identity, source, policy)
|
|
return nil
|
|
}
|
|
})
|
|
return queued, err
|
|
}
|
|
|
|
func (s *Service) resolveAdmissionAutoMergePolicy(
|
|
ctx context.Context,
|
|
identity *QueueSessionIdentity,
|
|
sessionID string,
|
|
) *AutoMergePolicy {
|
|
if identity == nil {
|
|
policy := s.globalAutoMergePolicy()
|
|
return &policy
|
|
}
|
|
policy, err := s.ResolveAutoMergePolicy(ctx, *identity)
|
|
if err == nil {
|
|
return &policy
|
|
}
|
|
s.logger.Warn("automatic merge policy unavailable; admitting separate message",
|
|
zap.String("session_id", sessionID),
|
|
zap.Error(err))
|
|
return nil
|
|
}
|
|
|
|
func (s *Service) autoMergeCandidate(
|
|
ctx context.Context,
|
|
identity *QueueSessionIdentity,
|
|
candidate *QueuedMessage,
|
|
claim *QueueAttachmentClaim,
|
|
policy *AutoMergePolicy,
|
|
) (*QueuedMessage, bool, error) {
|
|
if identity == nil {
|
|
return s.repo.AutoMergeCandidateIntoAbove(ctx, candidate)
|
|
}
|
|
if policyRepo, ok := s.repo.(autoMergeAdmissionRepository); ok && policy != nil {
|
|
return policyRepo.AutoMergeCandidateIntoAboveForSessionWithPolicy(ctx, *identity, candidate, claim, *policy)
|
|
}
|
|
if claim == nil {
|
|
return s.repo.AutoMergeCandidateIntoAboveForSession(ctx, *identity, candidate)
|
|
}
|
|
claimRepo, ok := s.repo.(queueAttachmentRepository)
|
|
if !ok {
|
|
return nil, false, errors.New("transactional attachment queue repository is unavailable")
|
|
}
|
|
return claimRepo.AutoMergeCandidateIntoAboveForSessionWithClaim(ctx, *identity, candidate, *claim)
|
|
}
|
|
|
|
// admitQueueFullMessage folds a candidate into the tail or retries insertion
|
|
// when the fold observes capacity freed by a concurrent drain.
|
|
func (s *Service) admitQueueFullMessage(
|
|
ctx context.Context,
|
|
identity *QueueSessionIdentity,
|
|
queueFullErr error,
|
|
sessionID, taskID, content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
claim *QueueAttachmentClaim,
|
|
maxPerSession int,
|
|
policy *AutoMergePolicy,
|
|
) (*QueuedMessage, *QueuedMessage, error) {
|
|
candidate := &QueuedMessage{
|
|
SessionID: sessionID,
|
|
TaskID: taskID,
|
|
Content: content,
|
|
Model: model,
|
|
PlanMode: planMode,
|
|
Attachments: attachments,
|
|
Metadata: copyMessageMetadata(metadata, 0),
|
|
QueuedAt: time.Now().UTC(),
|
|
QueuedBy: userID,
|
|
}
|
|
merged, didMerge, err := s.autoMergeCandidate(ctx, identity, candidate, claim, policy)
|
|
if errors.Is(err, ErrAutoMergePolicyChanged) {
|
|
return nil, nil, err
|
|
}
|
|
if err != nil {
|
|
if errors.Is(err, ErrTaskInactive) {
|
|
return nil, nil, err
|
|
}
|
|
s.logger.Error("automatic merge into full queue failed; preserving queue full rejection",
|
|
zap.String("session_id", sessionID),
|
|
zap.Error(err))
|
|
return nil, nil, queueFullErr
|
|
}
|
|
if didMerge && merged != nil {
|
|
s.logger.Info("automatically merged queued entry into full tail",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("surviving_entry_id", merged.ID))
|
|
return nil, merged, nil
|
|
}
|
|
source, err := s.insertQueueMessageWithMetadata(
|
|
ctx, identity, sessionID, taskID, content, model, userID, planMode,
|
|
attachments, metadata, claim, maxPerSession, policy,
|
|
)
|
|
return source, nil, err
|
|
}
|
|
|
|
// queueMessageWithMetadataSeparate is used by retry paths whose existing
|
|
// contract must not gain admission-time automatic behavior.
|
|
func (s *Service) queueMessageWithMetadataSeparate(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, maxPerSession int) (*QueuedMessage, error) {
|
|
var queued *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
queued, err = s.insertQueueMessageWithMetadata(
|
|
admittedCtx, nil, sessionID, taskID, content, model, userID, planMode,
|
|
attachments, metadata, nil, maxPerSession, nil,
|
|
)
|
|
return err
|
|
})
|
|
return queued, err
|
|
}
|
|
|
|
func (s *Service) finalizeAutoMerge(
|
|
ctx context.Context,
|
|
identity *QueueSessionIdentity,
|
|
source *QueuedMessage,
|
|
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
|
|
if identity == nil {
|
|
merged, didMerge, err = s.repo.AutoMergeIntoAbove(ctx, source.SessionID, source.ID)
|
|
} else if policyRepo, ok := s.repo.(autoMergeAdmissionRepository); ok {
|
|
merged, didMerge, err = policyRepo.AutoMergeIntoAboveForSessionWithPolicy(ctx, *identity, source.ID, *policy)
|
|
} else {
|
|
merged, didMerge, err = s.repo.AutoMergeIntoAboveForSession(ctx, *identity, source.ID)
|
|
}
|
|
if errors.Is(err, ErrAutoMergePolicyChanged) && identity != nil {
|
|
policy = s.resolveAdmissionAutoMergePolicy(ctx, identity, source.SessionID)
|
|
continue
|
|
}
|
|
if err != nil {
|
|
s.logger.Error("automatic queue merge failed; preserving separate admission",
|
|
zap.String("session_id", source.SessionID),
|
|
zap.String("source_entry_id", source.ID),
|
|
zap.Error(err))
|
|
return source
|
|
}
|
|
if !didMerge || merged == nil {
|
|
return source
|
|
}
|
|
s.logger.Info("automatically merged queued entry into tail",
|
|
zap.String("session_id", source.SessionID),
|
|
zap.String("source_entry_id", source.ID),
|
|
zap.String("surviving_entry_id", merged.ID))
|
|
return merged
|
|
}
|
|
return source
|
|
}
|
|
|
|
func (s *Service) insertQueueMessage(
|
|
ctx context.Context,
|
|
identity *QueueSessionIdentity,
|
|
msg *QueuedMessage,
|
|
claim *QueueAttachmentClaim,
|
|
maxPerSession int,
|
|
policy *AutoMergePolicy,
|
|
) error {
|
|
if identity == nil {
|
|
return s.repo.Insert(ctx, msg, maxPerSession)
|
|
}
|
|
if policyRepo, ok := s.repo.(autoMergeAdmissionRepository); ok && policy != nil {
|
|
return policyRepo.InsertForSessionWithPolicy(ctx, *identity, msg, claim, maxPerSession, *policy)
|
|
}
|
|
if claim == nil {
|
|
return s.repo.InsertForSession(ctx, *identity, msg, maxPerSession)
|
|
}
|
|
claimRepo, ok := s.repo.(queueAttachmentRepository)
|
|
if !ok {
|
|
return errors.New("transactional attachment queue repository is unavailable")
|
|
}
|
|
return claimRepo.InsertForSessionWithClaim(ctx, *identity, msg, *claim, maxPerSession)
|
|
}
|
|
|
|
// insertQueueMessageWithMetadata inserts a message with metadata under the per-session admission lock.
|
|
func (s *Service) insertQueueMessageWithMetadata(ctx context.Context, identity *QueueSessionIdentity, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, claim *QueueAttachmentClaim, maxPerSession int, policy *AutoMergePolicy) (*QueuedMessage, error) {
|
|
metadataCopy := copyMessageMetadata(metadata, 0)
|
|
msg := &QueuedMessage{
|
|
SessionID: sessionID,
|
|
TaskID: taskID,
|
|
Content: content,
|
|
Model: model,
|
|
PlanMode: planMode,
|
|
Attachments: attachments,
|
|
Metadata: metadataCopy,
|
|
QueuedBy: userID,
|
|
}
|
|
err := s.insertQueueMessage(ctx, identity, msg, claim, maxPerSession, policy)
|
|
if err != nil {
|
|
if errors.Is(err, ErrQueueFull) {
|
|
s.logger.Info("queue full",
|
|
zap.String("session_id", sessionID),
|
|
zap.Int("max", maxPerSession))
|
|
}
|
|
return nil, err
|
|
}
|
|
s.logger.Info("message queued",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("task_id", taskID),
|
|
zap.String("entry_id", msg.ID),
|
|
zap.Int64("position", msg.Position),
|
|
zap.Int("content_length", len(content)))
|
|
return msg, nil
|
|
}
|
|
|
|
// RestoreMessage restores a dequeued message at its original FIFO position
|
|
// after a delivery failure. It preserves the entry's identity, position,
|
|
// queued_at time, and queued_by ownership.
|
|
func (s *Service) RestoreMessage(ctx context.Context, msg *QueuedMessage) (*QueuedMessage, error) {
|
|
if msg == nil {
|
|
return nil, errors.New("queued message is nil")
|
|
}
|
|
var restored *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, msg.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
restored, err = s.restoreMessage(admittedCtx, msg)
|
|
return err
|
|
})
|
|
return restored, err
|
|
}
|
|
|
|
func (s *Service) RestoreMessageForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
msg *QueuedMessage,
|
|
) (*QueuedMessage, error) {
|
|
if msg == nil {
|
|
return nil, errors.New("queued message is nil")
|
|
}
|
|
var restored *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
candidate := *msg
|
|
candidate.Metadata = copyMessageMetadata(msg.Metadata, 0)
|
|
if err := s.validateSessionIdentity(admittedCtx, identity); err != nil {
|
|
return err
|
|
}
|
|
if err := s.repo.RestoreForSession(admittedCtx, identity, &candidate, 0); err != nil {
|
|
return err
|
|
}
|
|
restored = &candidate
|
|
return nil
|
|
})
|
|
return restored, err
|
|
}
|
|
|
|
// restoreMessage is the admission-checked core of RestoreMessage.
|
|
func (s *Service) restoreMessage(ctx context.Context, msg *QueuedMessage) (*QueuedMessage, error) {
|
|
restored := *msg
|
|
restored.Metadata = copyMessageMetadata(msg.Metadata, 0)
|
|
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),
|
|
zap.String("entry_id", restored.ID),
|
|
zap.Int64("position", restored.Position))
|
|
return &restored, nil
|
|
}
|
|
|
|
// QueueMessageWithCoalesceKey replaces an existing pending entry with the same
|
|
// coalesce key, session, and queued_by value. When no matching entry exists it
|
|
// inserts a new tail entry if allowInsert is true; otherwise ErrEntryNotFound is
|
|
// returned. The returned bool is true when an existing entry was replaced.
|
|
func (s *Service) QueueMessageWithCoalesceKey(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, coalesceKey string, allowInsert bool) (*QueuedMessage, bool, error) {
|
|
return s.queueMessageWithCoalesceKey(
|
|
ctx, sessionID, taskID, content, model, userID, planMode, attachments, metadata,
|
|
coalesceKey, allowInsert, s.MaxPerSession(),
|
|
)
|
|
}
|
|
|
|
func (s *Service) QueueMessageWithCoalesceKeyForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
coalesceKey string,
|
|
allowInsert bool,
|
|
) (*QueuedMessage, bool, error) {
|
|
var queued *QueuedMessage
|
|
var replaced bool
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
metadataCopy := copyMessageMetadata(metadata, 1)
|
|
metadataCopy[MetadataCoalesceKey] = coalesceKey
|
|
candidate := &QueuedMessage{
|
|
SessionID: identity.SessionID, TaskID: identity.TaskID,
|
|
Content: content, Model: model, PlanMode: planMode,
|
|
Attachments: attachments, Metadata: metadataCopy, QueuedBy: userID,
|
|
}
|
|
var err error
|
|
queued, replaced, err = s.repo.InsertOrReplaceByCoalesceKeyForSession(
|
|
admittedCtx, identity, candidate, coalesceKey, s.MaxPerSession(), allowInsert,
|
|
)
|
|
return err
|
|
})
|
|
return queued, replaced, err
|
|
}
|
|
|
|
// queueMessageWithCoalesceKey is the admission-checked core of QueueMessageWithCoalesceKey.
|
|
func (s *Service) queueMessageWithCoalesceKey(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, coalesceKey string, allowInsert bool, maxPerSession int) (*QueuedMessage, bool, error) {
|
|
var queued *QueuedMessage
|
|
var replaced bool
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
queued, replaced, err = s.insertQueueMessageWithCoalesceKey(
|
|
admittedCtx, sessionID, taskID, content, model, userID, planMode, attachments, metadata,
|
|
coalesceKey, allowInsert, maxPerSession,
|
|
)
|
|
return err
|
|
})
|
|
return queued, replaced, err
|
|
}
|
|
|
|
// insertQueueMessageWithCoalesceKey inserts or coalesce-replaces an entry under the admission lock.
|
|
func (s *Service) insertQueueMessageWithCoalesceKey(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, coalesceKey string, allowInsert bool, maxPerSession int) (*QueuedMessage, bool, error) {
|
|
metadataCopy := copyMessageMetadata(metadata, 1)
|
|
metadataCopy[MetadataCoalesceKey] = coalesceKey
|
|
msg := &QueuedMessage{
|
|
SessionID: sessionID,
|
|
TaskID: taskID,
|
|
Content: content,
|
|
Model: model,
|
|
PlanMode: planMode,
|
|
Attachments: attachments,
|
|
Metadata: metadataCopy,
|
|
QueuedBy: userID,
|
|
}
|
|
queued, replaced, err := s.repo.InsertOrReplaceByCoalesceKey(ctx, msg, coalesceKey, maxPerSession, allowInsert)
|
|
if err != nil {
|
|
if errors.Is(err, ErrQueueFull) {
|
|
s.logger.Info("queue full",
|
|
zap.String("session_id", sessionID),
|
|
zap.Int("max", maxPerSession))
|
|
}
|
|
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),
|
|
zap.String("entry_id", queued.ID),
|
|
zap.String("coalesce_key", coalesceKey),
|
|
zap.Bool("replaced", replaced),
|
|
zap.Int64("position", queued.Position),
|
|
zap.Int("content_length", len(content)))
|
|
return queued, replaced, nil
|
|
}
|
|
|
|
// RequeueMessage retries work that was already admitted. It bypasses the
|
|
// current admission cap so lowering capacity while a message is in flight
|
|
// cannot discard that message. Coalesced retries retain replacement semantics.
|
|
func (s *Service) RequeueMessage(ctx context.Context, msg *QueuedMessage, queuedBy, coalesceKey string) (*QueuedMessage, bool, error) {
|
|
if msg == nil {
|
|
return nil, false, errors.New("queued message is nil")
|
|
}
|
|
if coalesceKey != "" {
|
|
return s.queueMessageWithCoalesceKey(
|
|
ctx, msg.SessionID, msg.TaskID, msg.Content, msg.Model, queuedBy, msg.PlanMode,
|
|
msg.Attachments, msg.Metadata, coalesceKey, true, 0,
|
|
)
|
|
}
|
|
queued, err := s.queueMessageWithMetadataSeparate(
|
|
ctx, msg.SessionID, msg.TaskID, msg.Content, msg.Model, queuedBy, msg.PlanMode,
|
|
msg.Attachments, msg.Metadata, 0,
|
|
)
|
|
return queued, false, err
|
|
}
|
|
|
|
// RequeueAtHead re-enqueues a user message at a position strictly lower
|
|
// than the session's current head. It is the FIFO-preserving requeue
|
|
// that Service.requeueMessage invokes when an entry was superseded by a
|
|
// newer dispatch before it could be claimed. See Repository
|
|
// .RequeuePreservingFIFO for the position arithmetic and position rebasing.
|
|
//
|
|
// This method is the bug fix entry point — caller code that wants
|
|
// user-message retry should call this instead of RequeueMessage so
|
|
// the original message beats any new arrival on a busy session.
|
|
func (s *Service) RequeueAtHead(ctx context.Context, msg *QueuedMessage) error {
|
|
if msg == nil {
|
|
return errors.New("queued message is nil")
|
|
}
|
|
return s.WithSessionAdmission(ctx, msg.SessionID, func(admittedCtx context.Context) error {
|
|
if err := s.repo.RequeuePreservingFIFO(admittedCtx, msg); err != nil {
|
|
return err
|
|
}
|
|
s.invalidateEditLease(msg.SessionID, msg.ID)
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (s *Service) RequeueAtHeadForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
msg *QueuedMessage,
|
|
) error {
|
|
if msg == nil {
|
|
return errors.New("queued message is nil")
|
|
}
|
|
return s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
if err := s.validateSessionIdentity(admittedCtx, identity); err != nil {
|
|
return err
|
|
}
|
|
return s.repo.RequeuePreservingFIFOForSession(admittedCtx, identity, msg)
|
|
})
|
|
}
|
|
|
|
// QueueLifecycleMessageWithCoalesceKey accepts a lifecycle entry only while
|
|
// its task remains active. accepted is false for a normal archive/delete win.
|
|
func (s *Service) QueueLifecycleMessageWithCoalesceKey(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, coalesceKey string, allowInsert bool) (*QueuedMessage, bool, bool, error) {
|
|
return s.queueLifecycleMessageWithCoalesceKey(ctx, nil, sessionID, taskID, content, model, userID, planMode, attachments, metadata, coalesceKey, allowInsert, false, nil)
|
|
}
|
|
|
|
func (s *Service) QueueLifecycleMessageWithCoalesceKeyForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
coalesceKey string,
|
|
allowInsert bool,
|
|
) (*QueuedMessage, bool, bool, error) {
|
|
return s.queueLifecycleMessageWithCoalesceKey(
|
|
ctx, &identity, identity.SessionID, identity.TaskID, content, model, userID, planMode,
|
|
attachments, metadata, coalesceKey, allowInsert, false, nil,
|
|
)
|
|
}
|
|
|
|
// QueueLifecycleMessageWithCoalesceKeyAfterInsert runs afterInsert while the
|
|
// queue admission lock is held. If the callback fails, the exact queue
|
|
// snapshot from before admission is restored before the error is returned.
|
|
func (s *Service) QueueLifecycleMessageWithCoalesceKeyAfterInsert(
|
|
ctx context.Context,
|
|
sessionID, taskID, content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
coalesceKey string,
|
|
allowInsert bool,
|
|
afterInsert func(context.Context, *QueuedMessage, bool) error,
|
|
) (*QueuedMessage, bool, bool, error) {
|
|
return s.queueLifecycleMessageWithCoalesceKey(ctx, nil, sessionID, taskID, content, model, userID, planMode, attachments, metadata, coalesceKey, allowInsert, false, afterInsert)
|
|
}
|
|
|
|
func (s *Service) QueueLifecycleMessageWithCoalesceKeyForSessionAfterInsert(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
coalesceKey string,
|
|
allowInsert bool,
|
|
afterInsert func(context.Context, *QueuedMessage, bool) error,
|
|
) (*QueuedMessage, bool, bool, error) {
|
|
return s.queueLifecycleMessageWithCoalesceKey(
|
|
ctx, &identity, identity.SessionID, identity.TaskID, content, model, userID, planMode,
|
|
attachments, metadata, coalesceKey, allowInsert, false, afterInsert,
|
|
)
|
|
}
|
|
|
|
// RequeueLifecycleMessageWithCoalesceKey preserves the generation captured by
|
|
// an existing durable row. This prevents a stale retry from becoming a new
|
|
// prompt after the task was archived and then unarchived.
|
|
func (s *Service) RequeueLifecycleMessageWithCoalesceKey(ctx context.Context, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, coalesceKey string, allowInsert bool) (*QueuedMessage, bool, bool, error) {
|
|
return s.queueLifecycleMessageWithCoalesceKey(ctx, nil, sessionID, taskID, content, model, userID, planMode, attachments, metadata, coalesceKey, allowInsert, true, nil)
|
|
}
|
|
|
|
func (s *Service) RequeueLifecycleMessageWithCoalesceKeyForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
content, model, userID string,
|
|
planMode bool,
|
|
attachments []MessageAttachment,
|
|
metadata map[string]interface{},
|
|
coalesceKey string,
|
|
allowInsert bool,
|
|
) (*QueuedMessage, bool, bool, error) {
|
|
return s.queueLifecycleMessageWithCoalesceKey(
|
|
ctx, &identity, identity.SessionID, identity.TaskID, content, model, userID, planMode,
|
|
attachments, metadata, coalesceKey, allowInsert, true, nil,
|
|
)
|
|
}
|
|
|
|
// queueLifecycleMessageWithCoalesceKey is the lifecycle-guarded core of the lifecycle queue methods.
|
|
func (s *Service) queueLifecycleMessageWithCoalesceKey(ctx context.Context, identity *QueueSessionIdentity, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, coalesceKey string, allowInsert, isRetry bool, afterInsert func(context.Context, *QueuedMessage, bool) error) (*QueuedMessage, bool, bool, error) {
|
|
var queued *QueuedMessage
|
|
var replaced, accepted bool
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
if identity != nil {
|
|
if err := s.validateSessionIdentity(admittedCtx, *identity); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
var snapshot []QueuedMessage
|
|
var pendingMove *PendingMove
|
|
if afterInsert != nil {
|
|
var snapshotErr error
|
|
snapshot, snapshotErr = s.repo.ListBySession(admittedCtx, sessionID)
|
|
if snapshotErr != nil {
|
|
return snapshotErr
|
|
}
|
|
pendingMove, snapshotErr = s.repo.GetPendingMove(admittedCtx, sessionID)
|
|
if snapshotErr != nil {
|
|
return snapshotErr
|
|
}
|
|
}
|
|
var err error
|
|
queued, replaced, accepted, err = s.insertLifecycleMessageWithCoalesceKey(
|
|
admittedCtx, identity, sessionID, taskID, content, model, userID, planMode, attachments,
|
|
metadata, coalesceKey, allowInsert, isRetry,
|
|
)
|
|
if err != nil || !accepted || afterInsert == nil {
|
|
return err
|
|
}
|
|
if err := afterInsert(admittedCtx, queued, replaced); err != nil {
|
|
rollbackErr := s.restoreCoalescedAdmission(context.WithoutCancel(admittedCtx), sessionID, snapshot, pendingMove)
|
|
queued = nil
|
|
replaced = false
|
|
accepted = false
|
|
if rollbackErr != nil {
|
|
return fmt.Errorf("coalesced lifecycle admission callback failed: %w (queue rollback failed: %v)", err, rollbackErr)
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
})
|
|
return queued, replaced, accepted, err
|
|
}
|
|
|
|
func (s *Service) restoreCoalescedAdmission(ctx context.Context, sessionID string, entries []QueuedMessage, pendingMove *PendingMove) error {
|
|
if err := s.repo.ReplaceSession(ctx, sessionID, entries, pendingMove); err != nil {
|
|
s.logger.Error("failed to roll back lifecycle queue admission",
|
|
zap.String("session_id", sessionID), zap.Error(err))
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// insertLifecycleMessageWithCoalesceKey inserts or replaces a lifecycle entry under the admission lock.
|
|
func (s *Service) insertLifecycleMessageWithCoalesceKey(ctx context.Context, identity *QueueSessionIdentity, sessionID, taskID, content, model, userID string, planMode bool, attachments []MessageAttachment, metadata map[string]interface{}, coalesceKey string, allowInsert, isRetry bool) (*QueuedMessage, bool, bool, error) {
|
|
metadataCopy := clearReservedMetadata(metadata)
|
|
generation, err := s.repo.LifecycleGeneration(ctx, taskID)
|
|
if err != nil {
|
|
return nil, false, false, err
|
|
}
|
|
if isRetry {
|
|
// Rows written before generation metadata existed belong to generation
|
|
// zero. They may retry while no purge occurred, but a later archive
|
|
// advances generation and rejects them just like newer rows.
|
|
expected, ok := lifecycleGenerationFromMetadata(metadataCopy)
|
|
if !ok {
|
|
expected = 0
|
|
}
|
|
if expected != generation {
|
|
return nil, false, false, nil
|
|
}
|
|
}
|
|
metadataCopy[MetadataCoalesceKey] = coalesceKey
|
|
metadataCopy[MetadataLifecycleDurable] = true
|
|
metadataCopy[MetadataLifecycleGeneration] = generation
|
|
msg := &QueuedMessage{SessionID: sessionID, TaskID: taskID, Content: content, Model: model, PlanMode: planMode, Attachments: attachments, Metadata: metadataCopy, QueuedBy: userID}
|
|
maxPerSession := s.MaxPerSession()
|
|
if isRetry {
|
|
maxPerSession = 0
|
|
}
|
|
var queued *QueuedMessage
|
|
var replaced bool
|
|
if identity == nil {
|
|
queued, replaced, err = s.repo.InsertOrReplaceLifecycleByCoalesceKey(ctx, msg, coalesceKey, maxPerSession, allowInsert)
|
|
} else {
|
|
queued, replaced, err = s.repo.InsertOrReplaceLifecycleByCoalesceKeyForSession(ctx, *identity, msg, coalesceKey, maxPerSession, allowInsert)
|
|
}
|
|
if errors.Is(err, ErrTaskInactive) || errors.Is(err, ErrLifecycleCancelled) {
|
|
return nil, false, false, nil
|
|
}
|
|
if err != nil {
|
|
return nil, false, false, err
|
|
}
|
|
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) {
|
|
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.
|
|
func lifecycleGenerationFromMetadata(metadata map[string]interface{}) (int64, bool) {
|
|
if metadata == nil {
|
|
return 0, false
|
|
}
|
|
switch value := metadata[MetadataLifecycleGeneration].(type) {
|
|
case int64:
|
|
return value, true
|
|
case int:
|
|
return int64(value), true
|
|
case float64:
|
|
return int64(value), true
|
|
default:
|
|
return 0, false
|
|
}
|
|
}
|
|
|
|
// ReserveQueued atomically takes an ordinary head entry or reserves a durable
|
|
// lifecycle head entry. The admission lock keeps drains from observing an
|
|
// insert before its automatic-merge finalization completes. A reserved
|
|
// lifecycle row survives until acknowledged. Auto-run OFF leaves the head in
|
|
// place and reports no reserved entry.
|
|
func (s *Service) ReserveQueued(ctx context.Context, sessionID string) (*QueuedMessage, bool) {
|
|
msg, exists, _ := s.ReserveQueuedWithAutoRun(ctx, sessionID)
|
|
return msg, exists
|
|
}
|
|
|
|
// ReserveQueuedWithAutoRun also reports the policy decision. Nil/false/true
|
|
// means enabled but empty (or a logged storage error); nil/false/false means
|
|
// Auto-run is OFF.
|
|
func (s *Service) ReserveQueuedWithAutoRun(ctx context.Context, sessionID string) (*QueuedMessage, bool, bool) {
|
|
var msg *QueuedMessage
|
|
autoRun := true
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) 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
|
|
})
|
|
if err != nil {
|
|
s.logger.Error("reserve head failed",
|
|
zap.String("session_id", sessionID),
|
|
zap.Error(err))
|
|
return nil, false, true
|
|
}
|
|
return msg, msg != nil, autoRun
|
|
}
|
|
|
|
// ReserveQueuedWithAutoRunForSession reserves the FIFO head only for the exact session incarnation.
|
|
func (s *Service) ReserveQueuedWithAutoRunForSession(ctx context.Context, identity QueueSessionIdentity) (*QueuedMessage, bool, bool, error) {
|
|
var msg *QueuedMessage
|
|
var autoRun bool
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
msg, autoRun, err = s.repo.ReserveHeadIfAutoRunForSession(admittedCtx, identity)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return nil, false, true, err
|
|
}
|
|
return msg, msg != nil, autoRun, nil
|
|
}
|
|
|
|
// ReserveQueuedForDeliveryWithAutoRunForSession retains the FIFO head until
|
|
// explicit acknowledgement, regardless of producer type.
|
|
func (s *Service) ReserveQueuedForDeliveryWithAutoRunForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
) (*QueuedMessage, bool, bool, error) {
|
|
var msg *QueuedMessage
|
|
var autoRun bool
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
msg, autoRun, err = s.repo.ReserveHeadForDeliveryIfAutoRunForSession(admittedCtx, identity)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return nil, false, true, err
|
|
}
|
|
return msg, msg != nil, autoRun, nil
|
|
}
|
|
|
|
// SetAutoRun persists automatic-drain policy through the queue admission gate.
|
|
func (s *Service) SetAutoRun(ctx context.Context, sessionID string, enabled bool) error {
|
|
return s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
return s.repo.SetAutoRun(admittedCtx, sessionID, enabled)
|
|
})
|
|
}
|
|
|
|
// SetAutoRunForSession persists automatic-drain policy only for the exact session incarnation.
|
|
func (s *Service) SetAutoRunForSession(ctx context.Context, identity QueueSessionIdentity, enabled bool) error {
|
|
return s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
return s.repo.SetAutoRunForSession(admittedCtx, identity, enabled)
|
|
})
|
|
}
|
|
|
|
// PauseAutoRunIfPending atomically parks a visible backlog, if one exists.
|
|
func (s *Service) PauseAutoRunIfPending(ctx context.Context, sessionID string) (bool, error) {
|
|
var paused bool
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
paused, err = s.repo.PauseAutoRunIfPending(admittedCtx, sessionID)
|
|
return err
|
|
})
|
|
return paused, err
|
|
}
|
|
func (s *Service) PauseAutoRunIfPendingForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
) (bool, error) {
|
|
var paused bool
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
paused, err = s.repo.PauseAutoRunIfPendingForSession(admittedCtx, identity)
|
|
return err
|
|
})
|
|
return paused, err
|
|
}
|
|
|
|
// 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 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(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
reserved *QueuedMessage,
|
|
) error {
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
if err := s.validateSessionIdentity(admittedCtx, identity); err != nil {
|
|
return err
|
|
}
|
|
return s.repo.AcknowledgeByIDForSession(admittedCtx, identity, reserved)
|
|
})
|
|
if errors.Is(err, ErrEntryNotFound) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
|
|
// MarkDeliveryAttemptedForSession atomically crosses the at-most-once boundary
|
|
// for token-owned plan-comment receipts immediately before external I/O.
|
|
func (s *Service) MarkDeliveryAttemptedForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
messages []QueuedMessage,
|
|
) error {
|
|
return s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
return s.repo.MarkDeliveryAttemptedForSession(admittedCtx, identity, messages)
|
|
})
|
|
}
|
|
|
|
// ReleaseQueuedDeliveryForSession makes an unaccepted retained entry visible
|
|
// again for the same session incarnation.
|
|
func (s *Service) ReleaseQueuedDeliveryForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
reserved *QueuedMessage,
|
|
) error {
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
return s.repo.ReleaseDeliveryReservationForSession(admittedCtx, identity, reserved)
|
|
})
|
|
if errors.Is(err, ErrEntryNotFound) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
|
|
// DiscardLifecycleReservation removes a reservation still owned by a replaced
|
|
// session incarnation. The repository compares persisted ownership, so this
|
|
// cannot delete a reservation subsequently claimed by another incarnation.
|
|
func (s *Service) DiscardLifecycleReservation(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
reserved *QueuedMessage,
|
|
) error {
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
return s.repo.DiscardLifecycleReservation(admittedCtx, identity, reserved)
|
|
})
|
|
if errors.Is(err, ErrEntryNotFound) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
|
|
// IsCurrentLifecycleReservation verifies that msg is still the durable row
|
|
// accepted for the task's current lifecycle generation. A task archive/delete
|
|
// removes that row and advances the generation, so callers can discard a
|
|
// deferred reservation before it performs any prompt side effects.
|
|
func (s *Service) IsCurrentLifecycleReservation(ctx context.Context, msg *QueuedMessage) bool {
|
|
if msg == nil || !msg.IsDurableLifecycle() {
|
|
return false
|
|
}
|
|
expectedGeneration, ok := lifecycleGenerationFromMetadata(msg.Metadata)
|
|
if !ok {
|
|
expectedGeneration = 0
|
|
}
|
|
generation, err := s.repo.LifecycleGeneration(ctx, msg.TaskID)
|
|
if err != nil || generation != expectedGeneration {
|
|
return false
|
|
}
|
|
if msg.reservationIdentity.SessionIncarnationID != "" {
|
|
current, err := s.repo.ResolveSessionIdentity(
|
|
ctx,
|
|
msg.reservationIdentity.TaskID,
|
|
msg.reservationIdentity.SessionID,
|
|
)
|
|
if err != nil || current != msg.reservationIdentity {
|
|
return false
|
|
}
|
|
}
|
|
entries, err := s.repo.ListBySession(ctx, msg.SessionID)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
for i := range entries {
|
|
entry := &entries[i]
|
|
if entry.ID != msg.ID || entry.TaskID != msg.TaskID || !entry.IsDurableLifecycle() {
|
|
continue
|
|
}
|
|
if msg.reservationIdentity.SessionIncarnationID != "" &&
|
|
lifecycleReservationIncarnation(entry.Metadata) != msg.reservationIdentity.SessionIncarnationID {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// copyMessageMetadata clones metadata, optionally growing the map capacity.
|
|
func copyMessageMetadata(metadata map[string]interface{}, extraCapacity int) map[string]interface{} {
|
|
if len(metadata) == 0 && extraCapacity == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]interface{}, len(metadata)+extraCapacity)
|
|
for key, value := range metadata {
|
|
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)
|
|
}
|
|
|
|
// AppendContentForSession admits content only for the exact session incarnation.
|
|
func (s *Service) AppendContentForSession(ctx context.Context, identity QueueSessionIdentity, content, model, userID string, planMode bool, attachments []MessageAttachment) (*QueuedMessage, bool, error) {
|
|
var msg *QueuedMessage
|
|
var appended bool
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) 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
|
|
})
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
s.logger.Info("queue append-or-insert",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.String("entry_id", msg.ID),
|
|
zap.Bool("appended", appended))
|
|
return msg, appended, nil
|
|
}
|
|
|
|
// TakeQueued atomically removes and returns the head entry. Returns nil, false
|
|
func (s *Service) TakeQueued(ctx context.Context, sessionID string) (*QueuedMessage, bool) {
|
|
var msg *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) 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 {
|
|
s.logger.Error("take head failed",
|
|
zap.String("session_id", sessionID),
|
|
zap.Error(err))
|
|
return nil, false
|
|
}
|
|
if msg == nil {
|
|
return nil, false
|
|
}
|
|
s.logger.Info("message dequeued",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("entry_id", msg.ID))
|
|
return msg, true
|
|
}
|
|
|
|
// TakeQueuedIfAutoRun atomically checks the session policy and removes the
|
|
// FIFO head only while Auto-run is enabled. The admission gate serializes the
|
|
// policy read with every queue mutation made through this service.
|
|
func (s *Service) TakeQueuedIfAutoRun(ctx context.Context, sessionID string) (*QueuedMessage, bool) {
|
|
var msg *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
autoRun, err := s.repo.GetAutoRun(admittedCtx, sessionID)
|
|
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 {
|
|
s.logger.Error("take Auto-run head failed",
|
|
zap.String("session_id", sessionID),
|
|
zap.Error(err))
|
|
return nil, false
|
|
}
|
|
if msg == nil {
|
|
return nil, false
|
|
}
|
|
s.logger.Info("Auto-run message dequeued",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("entry_id", msg.ID))
|
|
return msg, true
|
|
}
|
|
|
|
// TakeQueuedEntry atomically removes and returns the entry identified by
|
|
// entryID, regardless of its FIFO position. Returns nil, false, nil when
|
|
// the entry no longer exists (already drained or removed by a concurrent
|
|
// path); returns a non-nil error on a genuine repository failure — the two
|
|
// are deliberately distinguished so InterruptForPeerMessage can treat "not
|
|
// found" (safe to fall back to a FIFO-head drain) differently from "the
|
|
// repository call itself failed" (unsafe to fall back: a transient error
|
|
// here says nothing about which entry is actually at the FIFO head, and
|
|
// falling back anyway risks dispatching the wrong message while still
|
|
// reporting "sent" for the parent's — see InterruptForPeerMessage). Used by
|
|
// InterruptForPeerMessage to dispatch the specific message that triggered
|
|
// the interrupt instead of whatever happens to be at the FIFO head.
|
|
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 {
|
|
s.logger.Error("take by id failed",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("entry_id", entryID),
|
|
zap.Error(err))
|
|
return nil, false, err
|
|
}
|
|
if msg == nil {
|
|
return nil, false, nil
|
|
}
|
|
s.logger.Info("message dequeued by id",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("entry_id", msg.ID))
|
|
return msg, true, nil
|
|
}
|
|
|
|
func (s *Service) TakeQueuedEntryForSession(
|
|
ctx context.Context,
|
|
identity QueueSessionIdentity,
|
|
entryID string,
|
|
) (*QueuedMessage, bool, error) {
|
|
var msg *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
msg, err = s.repo.TakeByIDForSession(admittedCtx, identity, entryID)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
s.logger.Error("take by id for session failed",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.String("entry_id", entryID),
|
|
zap.Error(err))
|
|
return nil, false, err
|
|
}
|
|
if msg == nil {
|
|
return nil, false, nil
|
|
}
|
|
s.logger.Info("message dequeued by id",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.String("entry_id", msg.ID))
|
|
return msg, true, nil
|
|
}
|
|
|
|
// UpdateMessage replaces the content (and optionally attachments) of a queued
|
|
// entry. The sessionID scope is mandatory — callers can't update an entry by
|
|
// guessing its UUID across sessions. Returns ErrEntryNotFound when the entry
|
|
// was already drained, the session doesn't own it, or the queuedBy guard
|
|
// rejects the caller.
|
|
func (s *Service) UpdateMessage(ctx context.Context, sessionID, entryID, content string, attachments []MessageAttachment, queuedBy string) error {
|
|
return s.UpdateMessageWithMetadata(ctx, sessionID, entryID, content, attachments, nil, queuedBy)
|
|
}
|
|
|
|
// GetEntry returns a pending queue entry scoped to its session. It is used by
|
|
// 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) {
|
|
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
|
|
}
|
|
return referenced, nil
|
|
}
|
|
|
|
// ClaimSendNow atomically claims the exact pending source snapshot for an
|
|
// interrupt-and-replace dispatch. The repository orders the retained sources
|
|
// by FIFO position and constructs the synthetic dispatch envelope before
|
|
// mutating any row, so aggregate validation failures or click-time edits leave
|
|
// the queue untouched.
|
|
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 {
|
|
return nil, err
|
|
}
|
|
s.logger.Info("claimed queued messages for send now",
|
|
zap.String("session_id", sessionID),
|
|
zap.Int("source_count", len(claim.Sources)))
|
|
return claim, nil
|
|
}
|
|
|
|
// ClaimSendNowForSession claims the exact pending snapshot only for the exact session incarnation.
|
|
func (s *Service) ClaimSendNowForSession(ctx context.Context, identity QueueSessionIdentity, expected []QueuedMessage) (*SendNowClaim, error) {
|
|
var claim *SendNowClaim
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
claim, err = s.repo.ClaimSendNowForSession(admittedCtx, identity, expected)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
s.logger.Info("claimed queued messages for send now",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.Int("source_count", len(claim.Sources)))
|
|
return claim, nil
|
|
}
|
|
|
|
// 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 claim == nil {
|
|
return s.repo.RestoreSendNowClaim(ctx, claim)
|
|
}
|
|
sessionID, err := sendNowClaimSessionID(claim)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
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
|
|
}
|
|
|
|
// AcknowledgeSendNowClaim removes durable lifecycle sources after the
|
|
// 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 claim == nil {
|
|
return s.repo.AcknowledgeSendNowClaim(ctx, claim)
|
|
}
|
|
sessionID, err := sendNowClaimSessionID(claim)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
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.IsDurableDelivery() {
|
|
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) {
|
|
snapshot, err := s.repo.Snapshot(ctx, identity)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for index := range snapshot.Entries {
|
|
if snapshot.Entries[index].ID == entryID {
|
|
entry := snapshot.Entries[index]
|
|
return &entry, nil
|
|
}
|
|
}
|
|
return nil, ErrEntryNotFound
|
|
}
|
|
|
|
// UpdateMessageWithMetadata atomically edits queue content and applies
|
|
// 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
|
|
}
|
|
s.logger.Info("queued entry updated",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("entry_id", entryID))
|
|
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 {
|
|
return s.repo.UpdateContentAndMetadataForSession(admittedCtx, identity, entryID, content, attachments, metadataUpdates, queuedBy)
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
s.logger.Info("queued entry updated",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.String("entry_id", entryID))
|
|
return nil
|
|
}
|
|
|
|
// UpdateMessageWithMetadataForSessionWithClaim claims newly referenced staged
|
|
// attachments and updates the exact queue row atomically.
|
|
func (s *Service) UpdateMessageWithMetadataForSessionWithClaim(ctx context.Context, identity QueueSessionIdentity, entryID, content string, attachments []MessageAttachment, metadataUpdates map[string]interface{}, queuedBy string, claim QueueAttachmentClaim) error {
|
|
if err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
claimRepo, ok := s.repo.(queueAttachmentRepository)
|
|
if !ok {
|
|
return errors.New("transactional attachment queue repository is unavailable")
|
|
}
|
|
return claimRepo.UpdateContentAndMetadataForSessionWithClaim(admittedCtx, identity, entryID, content, attachments, metadataUpdates, queuedBy, claim)
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
s.logger.Info("queued entry updated",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.String("entry_id", entryID))
|
|
return nil
|
|
}
|
|
|
|
// RemoveEntry deletes a single entry. The sessionID scope is mandatory — see
|
|
// 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 {
|
|
_, err := s.RemoveEntryWithEntry(ctx, sessionID, entryID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
s.logger.Info("queued entry removed",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("entry_id", entryID))
|
|
return nil
|
|
}
|
|
|
|
// RemoveEntryForSession deletes an entry only for the exact session incarnation
|
|
// and returns the removed entry plus every retained row from the same transaction.
|
|
func (s *Service) RemoveEntryForSession(ctx context.Context, identity QueueSessionIdentity, entryID string) (*QueueRemovalResult, error) {
|
|
var result *QueueRemovalResult
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
result, err = s.repo.DeleteByIDForSession(admittedCtx, identity, entryID)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
s.logger.Info("queued entry removed",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.String("entry_id", entryID))
|
|
return result, nil
|
|
}
|
|
|
|
// MergeIntoAbove folds the entry identified by entryID into the entry directly
|
|
// above it within the same session. See Repository.MergeIntoAbove for the merge
|
|
// rules and error mapping (ErrEntryNotFound / ErrNoMergeTarget).
|
|
func (s *Service) MergeIntoAbove(ctx context.Context, sessionID, entryID, queuedBy string) (*QueuedMessage, error) {
|
|
if !s.MergeEnabled() {
|
|
return nil, ErrMergeDisabled
|
|
}
|
|
var merged *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, sessionID, func(admittedCtx context.Context) error {
|
|
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
|
|
}
|
|
s.logger.Info("queued entry merged into entry above",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("entry_id", merged.ID))
|
|
return merged, nil
|
|
}
|
|
|
|
// MergeIntoAboveForSession merges an entry only for the exact session incarnation.
|
|
func (s *Service) MergeIntoAboveForSession(ctx context.Context, identity QueueSessionIdentity, entryID, queuedBy string) (*QueuedMessage, error) {
|
|
if !s.MergeEnabled() {
|
|
return nil, ErrMergeDisabled
|
|
}
|
|
var merged *QueuedMessage
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
merged, err = s.repo.MergeIntoAboveForSession(admittedCtx, identity, entryID, queuedBy)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
s.logger.Info("queued entry merged into entry above",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.String("entry_id", merged.ID))
|
|
return merged, nil
|
|
}
|
|
|
|
// ReorderEntries rewrites the pending FIFO order of a session's queue to
|
|
// match orderedIDs. Returns ErrQueueChanged when the submitted order is empty,
|
|
// contains duplicates, or does not match the current visible pending set (a
|
|
// drain/remove/merge raced the request); callers refetch the authoritative
|
|
// queue.
|
|
func (s *Service) ReorderEntries(ctx context.Context, sessionID string, orderedIDs []string) error {
|
|
if err := validateReorderInput(orderedIDs); err != nil {
|
|
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
|
|
}
|
|
s.logger.Info("queued entries reordered",
|
|
zap.String("session_id", sessionID),
|
|
zap.Int("entry_count", len(orderedIDs)))
|
|
return nil
|
|
}
|
|
|
|
// ReorderEntriesForSession reorders pending entries only for the exact session incarnation.
|
|
func (s *Service) ReorderEntriesForSession(ctx context.Context, identity QueueSessionIdentity, orderedIDs []string) error {
|
|
if err := validateReorderInput(orderedIDs); err != nil {
|
|
return err
|
|
}
|
|
if err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
return s.repo.ReorderEntriesForSession(admittedCtx, identity, orderedIDs)
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
s.logger.Info("queued entries reordered",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.Int("entry_count", len(orderedIDs)))
|
|
return nil
|
|
}
|
|
|
|
// CancelAll clears every queued entry for a session. Returns the number of
|
|
// rows removed.
|
|
func (s *Service) CancelAll(ctx context.Context, sessionID string) (int, error) {
|
|
var n int
|
|
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 {
|
|
return 0, err
|
|
}
|
|
s.logger.Info("queue cancel-all",
|
|
zap.String("session_id", sessionID),
|
|
zap.Int("removed", n))
|
|
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) {
|
|
var result *QueueRemovalResult
|
|
err := s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
var err error
|
|
result, err = s.repo.DeleteAllBySessionForIdentity(admittedCtx, identity)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
s.logger.Info("queue cancel-all",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.Int("removed", len(result.Removed)))
|
|
return result, nil
|
|
}
|
|
|
|
// PurgeDeletedSession removes all ephemeral queue state after the authoritative
|
|
// task-session row is gone. It never matches a replacement incarnation.
|
|
func (s *Service) PurgeDeletedSession(ctx context.Context, identity QueueSessionIdentity) error {
|
|
purger, ok := s.repo.(deletedSessionRepository)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
return s.WithSessionAdmission(ctx, identity.SessionID, func(admittedCtx context.Context) error {
|
|
return purger.PurgeDeletedSession(admittedCtx, identity)
|
|
})
|
|
}
|
|
|
|
// GetStatus returns the full pending list and capacity info for a session.
|
|
func (s *Service) GetStatus(ctx context.Context, sessionID string) *QueueStatus {
|
|
maxPerSession := s.MaxPerSession()
|
|
mergeEnabled := s.MergeEnabled()
|
|
status := &QueueStatus{
|
|
Entries: []QueuedMessage{},
|
|
Count: 0,
|
|
Max: maxPerSession,
|
|
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 {
|
|
// Unattempted comment receipts remain queryable during lost-response reconciliation.
|
|
if entry.IsReservedInFlight() &&
|
|
(!entry.IsDurablePlanComment() || entry.IsDeliveryAttempted()) {
|
|
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.
|
|
func (s *Service) Snapshot(ctx context.Context, identity QueueSessionIdentity) (*QueueStatus, error) {
|
|
snapshot, err := s.repo.Snapshot(ctx, identity)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
pending := make([]QueuedMessage, 0, len(snapshot.Entries))
|
|
for _, entry := range snapshot.Entries {
|
|
if !entry.IsReservedInFlight() ||
|
|
(entry.IsDurablePlanComment() && !entry.IsDeliveryAttempted()) {
|
|
pending = append(pending, entry)
|
|
}
|
|
}
|
|
status := &QueueStatus{
|
|
Entries: pending,
|
|
Count: len(pending),
|
|
Max: s.MaxPerSession(),
|
|
TaskID: identity.TaskID,
|
|
SessionID: identity.SessionID,
|
|
SessionIncarnationID: identity.SessionIncarnationID,
|
|
StatusEpoch: s.statusEpoch,
|
|
StatusGeneration: snapshot.StatusGeneration,
|
|
AutoRun: snapshot.AutoRun,
|
|
MergeEnabled: s.MergeEnabled(),
|
|
}
|
|
if snapshot.AutoMergeError != nil {
|
|
s.logger.Warn("automatic merge policy unavailable while reading queue status",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.Error(snapshot.AutoMergeError))
|
|
return status, nil
|
|
}
|
|
var policy AutoMergePolicy
|
|
if snapshot.AutoMergeOverride != nil {
|
|
policy = AutoMergePolicy{
|
|
Enabled: snapshot.AutoMergeOverride.Enabled, Source: AutoMergeSourceSession,
|
|
Revision: snapshot.AutoMergeOverride.Revision,
|
|
}
|
|
} else {
|
|
policy, err = s.loadGlobalAutoMergePolicy(ctx)
|
|
if err != nil {
|
|
s.logger.Warn("global automatic merge policy unavailable while reading queue status",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.Error(err))
|
|
return status, nil
|
|
}
|
|
}
|
|
status.AutoMergeAvailable = true
|
|
status.AutoMergeEnabled = &policy.Enabled
|
|
status.AutoMergeSource = policy.Source
|
|
status.AutoMergeRevision = &policy.Revision
|
|
return status, nil
|
|
}
|
|
|
|
// CountPendingByTaskIDs returns the pending prompt count per task, keyed by
|
|
// task_id, for every requested task. Reserved in-flight lifecycle rows are
|
|
// excluded, matching GetStatus. Used by task-list assembly and the status
|
|
// summary projector to render per-task queued-prompt badges.
|
|
func (s *Service) CountPendingByTaskIDs(ctx context.Context, taskIDs []string) (map[string]int, error) {
|
|
counts, err := s.repo.CountPendingByTaskIDs(ctx, taskIDs)
|
|
if err != nil {
|
|
s.logger.Error("count pending by task ids failed",
|
|
zap.Int("task_count", len(taskIDs)),
|
|
zap.Error(err))
|
|
return nil, err
|
|
}
|
|
return counts, nil
|
|
}
|
|
|
|
// CountPendingByTask returns the pending prompt count for one task.
|
|
func (s *Service) CountPendingByTask(ctx context.Context, taskID string) (int, error) {
|
|
counts, err := s.CountPendingByTaskIDs(ctx, []string{taskID})
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return counts[taskID], nil
|
|
}
|
|
|
|
// SnapshotSession returns the complete persisted queue state for rollback.
|
|
// 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) {
|
|
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)
|
|
}
|
|
repo, ok := s.repo.(ownedSessionTransferRepository)
|
|
if !ok {
|
|
return errors.New("owned session transfer unavailable")
|
|
}
|
|
return repo.transferSessionOwned(ctx, oldSessionID, newSessionID, operationID)
|
|
}
|
|
|
|
// SnapshotSessionForIdentity captures queue and deferred-move state for one exact incarnation.
|
|
func (s *Service) SnapshotSessionForIdentity(ctx context.Context, identity QueueSessionIdentity) ([]QueuedMessage, *PendingMove, error) {
|
|
snapshot, err := s.repo.Snapshot(ctx, identity)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("snapshot session queue: %w", err)
|
|
}
|
|
return snapshot.Entries, snapshot.PendingMove, nil
|
|
}
|
|
|
|
// TransferSession moves any queued messages and pending move from one session
|
|
// to another. Used by workflow session switches. Returns an error so the
|
|
// caller can fail closed instead of silently leaving entries orphaned on the
|
|
// 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 {
|
|
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),
|
|
zap.Error(err))
|
|
return err
|
|
}
|
|
s.logger.Info("transferred queue between sessions",
|
|
zap.String("from_session_id", oldSessionID),
|
|
zap.String("to_session_id", newSessionID))
|
|
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 {
|
|
return err
|
|
}
|
|
if current != identity {
|
|
return ErrSessionIdentityMismatch
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// TransferSessionIdentities moves queue state between exact live incarnations of the same task.
|
|
func (s *Service) TransferSessionIdentities(ctx context.Context, source, destination QueueSessionIdentity) error {
|
|
if source.TaskID != destination.TaskID {
|
|
return ErrSessionIdentityMismatch
|
|
}
|
|
if err := s.validateSessionIdentity(ctx, source); err != nil {
|
|
return err
|
|
}
|
|
if err := s.validateSessionIdentity(ctx, destination); err != nil {
|
|
return err
|
|
}
|
|
if err := s.repo.TransferSessionIdentities(ctx, source, destination); err != nil {
|
|
s.logger.Error("transfer session failed",
|
|
zap.String("from_session_id", source.SessionID),
|
|
zap.String("to_session_id", destination.SessionID),
|
|
zap.Error(err))
|
|
return err
|
|
}
|
|
s.logger.Info("transferred queue between sessions",
|
|
zap.String("from_session_id", source.SessionID),
|
|
zap.String("to_session_id", destination.SessionID))
|
|
return nil
|
|
}
|
|
|
|
// 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 {
|
|
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))
|
|
return err
|
|
}
|
|
s.logger.Info("restored session queue",
|
|
zap.String("session_id", sessionID),
|
|
zap.Int("entries", len(entries)),
|
|
zap.Bool("pending_move", pendingMove != nil))
|
|
return nil
|
|
}
|
|
|
|
// RestoreSessionForIdentity replaces queue state only for the captured session incarnation.
|
|
func (s *Service) RestoreSessionForIdentity(ctx context.Context, identity QueueSessionIdentity, entries []QueuedMessage, pendingMove *PendingMove) error {
|
|
if err := s.validateSessionIdentity(ctx, identity); err != nil {
|
|
return err
|
|
}
|
|
if err := s.repo.ReplaceSessionForIdentity(ctx, identity, entries, pendingMove); err != nil {
|
|
s.logger.Error("restore session queue failed",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.Error(err))
|
|
return err
|
|
}
|
|
s.logger.Info("restored session queue",
|
|
zap.String("session_id", identity.SessionID),
|
|
zap.Int("entries", len(entries)),
|
|
zap.Bool("pending_move", pendingMove != nil))
|
|
return nil
|
|
}
|
|
|
|
// 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.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))
|
|
return err
|
|
}
|
|
s.logger.Info("pending move recorded",
|
|
zap.String("session_id", sessionID),
|
|
zap.String("task_id", move.TaskID),
|
|
zap.String("workflow_step_id", move.WorkflowStepID))
|
|
return nil
|
|
}
|
|
|
|
// GetPendingMoveWithError retrieves the pending move for a session without
|
|
// 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) {
|
|
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
|
|
}
|
|
if move == nil {
|
|
return nil, false, nil
|
|
}
|
|
return move, true, nil
|
|
}
|
|
|
|
// GetPendingMove retrieves the pending move for a session without removing it.
|
|
func (s *Service) GetPendingMove(ctx context.Context, sessionID string) (*PendingMove, bool) {
|
|
move, exists, err := s.GetPendingMoveWithError(ctx, sessionID)
|
|
if err != nil {
|
|
s.logger.Error("get pending move failed",
|
|
zap.String("session_id", sessionID),
|
|
zap.Error(err))
|
|
return nil, false
|
|
}
|
|
return move, exists
|
|
}
|
|
|
|
// TakePendingMove retrieves and removes the pending move for a session.
|
|
func (s *Service) TakePendingMove(ctx context.Context, sessionID string) (*PendingMove, bool) {
|
|
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),
|
|
zap.Error(err))
|
|
return nil, false
|
|
}
|
|
if move == nil {
|
|
return nil, false
|
|
}
|
|
return move, true
|
|
}
|