bkinnightskytw
/
agent-orchestrator
mirror of https://github.com/ComposioHQ/agent-orchestrator.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: harden notification delivery leases

This commit is contained in:
whoisasx 2026-06-01 03:01:05 +05:30
parent d4622fe223
commit 041c8c8f7f
7 changed files with 104 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
backend/internal/notification/delivery.go
View File

@ -4,12 +4,15 @@ import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"time"
"github.com/aoagents/agent-orchestrator/backend/internal/domain"
)
var ErrDeliveryUpdateConflict = errors.New("notification delivery update conflict")
const (
RouteDashboard = "dashboard"
RouteDesktop = "desktop"

8
backend/internal/notification/dispatcher_test.go
View File

@ -63,16 +63,16 @@ func (f *fakeRuntimeStore) ReleaseExpiredDeliveryLeases(context.Context, time.Ti
f.releases++
return 0, nil
}
func (f *fakeRuntimeStore) MarkDeliverySent(context.Context, string, string, time.Time) error {
func (f *fakeRuntimeStore) MarkDeliverySent(context.Context, string, string, string, time.Time) error {
return nil
}
func (f *fakeRuntimeStore) MarkDeliveryRetry(_ context.Context, _ string, _ string, _ string, next time.Time) error {
func (f *fakeRuntimeStore) MarkDeliveryRetry(_ context.Context, _ string, _ string, _ string, _ string, next time.Time, _ time.Time) error {
f.mu.Lock()
defer f.mu.Unlock()
f.retryNext = next
return nil
}
func (f *fakeRuntimeStore) MarkDeliveryFailed(context.Context, string, string, string, time.Time) error {
func (f *fakeRuntimeStore) MarkDeliveryFailed(context.Context, string, string, string, string, time.Time) error {
return nil
}
func (f *fakeRuntimeStore) MarkDeliverySkipped(context.Context, string, string, time.Time) error {
@ -142,7 +142,7 @@ func TestMarkDeliveryErrorUsesAttemptAwareBackoff(t *testing.T) {
mgr := NewManager(store, StaticSettings(cfg), discardLogger())
mgr.clock = func() time.Time { return now }
if err := mgr.MarkDeliveryError(context.Background(), "del_1", "timeout", "timed out"); err != nil {
if err := mgr.MarkDeliveryError(context.Background(), "del_1", "electron", "timeout", "timed out"); err != nil {
t.Fatal(err)
}
store.mu.Lock()

10
backend/internal/notification/manager.go
View File

@ -130,11 +130,11 @@ func (m *Manager) ClaimDesktopDeliveries(ctx context.Context, owner string, limi
return m.store.ClaimDueDeliveries(ctx, SinkAOApp, owner, m.clock().UTC(), limit, policy.LeaseTTL)
}
func (m *Manager) MarkDeliverySent(ctx context.Context, id, externalID string) error {
return m.store.MarkDeliverySent(ctx, id, externalID, m.clock().UTC())
func (m *Manager) MarkDeliverySent(ctx context.Context, id, owner, externalID string) error {
return m.store.MarkDeliverySent(ctx, id, owner, externalID, m.clock().UTC())
}
func (m *Manager) MarkDeliveryError(ctx context.Context, id, code, message string) error {
func (m *Manager) MarkDeliveryError(ctx context.Context, id, owner, code, message string) error {
settings := m.settings.Settings(ctx)
policy := RetryPolicyFromConfig(settings.Retry)
now := m.clock().UTC()
@ -143,7 +143,7 @@ func (m *Manager) MarkDeliveryError(ctx context.Context, id, code, message strin
// fetch the current delivery attempts and provide the attempt-aware next
// retry timestamp for retry_wait rows.
if ClassifyError(code) == ErrorPermanent {
return m.store.MarkDeliveryFailed(ctx, id, code, message, now)
return m.store.MarkDeliveryFailed(ctx, id, owner, code, message, now)
}
row, ok, err := m.store.GetDelivery(ctx, id)
if err != nil {
@ -153,5 +153,5 @@ func (m *Manager) MarkDeliveryError(ctx context.Context, id, code, message strin
return fmt.Errorf("notification delivery %s not found", id)
}
nextAttemptNo := row.Attempts + 1
return m.store.MarkDeliveryRetry(ctx, id, code, message, policy.NextAttemptAt(now, nextAttemptNo))
return m.store.MarkDeliveryRetry(ctx, id, owner, code, message, policy.NextAttemptAt(now, nextAttemptNo), now)
}

6
backend/internal/notification/store.go
View File

@ -18,8 +18,8 @@ type Store interface {
EnqueueDelivery(ctx context.Context, row DeliveryRow) (DeliveryRow, bool, error)
ClaimDueDeliveries(ctx context.Context, sink string, owner string, now time.Time, limit int, lease time.Duration) ([]DeliveryRow, error)
ReleaseExpiredDeliveryLeases(ctx context.Context, now time.Time) (int, error)
MarkDeliverySent(ctx context.Context, id string, externalID string, at time.Time) error
MarkDeliveryRetry(ctx context.Context, id string, errCode string, errMessage string, next time.Time) error
MarkDeliveryFailed(ctx context.Context, id string, errCode string, errMessage string, at time.Time) error
MarkDeliverySent(ctx context.Context, id string, owner string, externalID string, at time.Time) error
MarkDeliveryRetry(ctx context.Context, id string, owner string, errCode string, errMessage string, next time.Time, at time.Time) error
MarkDeliveryFailed(ctx context.Context, id string, owner string, errCode string, errMessage string, at time.Time) error
MarkDeliverySkipped(ctx context.Context, id string, reason string, at time.Time) error
}

1
backend/internal/storage/sqlite/migrations/0003_notification_deliveries.sql
View File

@ -1,6 +1,7 @@
-- +goose Up
-- +goose StatementBegin
ALTER TABLE notifications ADD COLUMN routed_at TIMESTAMP;
UPDATE notifications SET routed_at = updated_at WHERE routed_at IS NULL;
CREATE INDEX idx_notifications_unrouted
ON notifications(seq)
WHERE routed_at IS NULL;

44
backend/internal/storage/sqlite/notification_delivery_store.go
View File

@ -18,7 +18,10 @@ const deliveryColumns = `id, notification_id, notification_seq, project_id, sess
last_error_code, last_error, external_id,
created_at, updated_at, delivered_at`
const defaultDeliveryLimit = 100
const (
defaultDeliveryLimit = 100
defaultDeliveryMaxAttempts = 5 // mirrors notification_deliveries.max_attempts schema default
)
type DeliveryFilter struct {
NotificationID string
@ -59,7 +62,7 @@ func (s *Store) EnqueueDelivery(ctx context.Context, row notification.DeliveryRo
if now.IsZero() {
now = time.Now().UTC()
}
row, err := notification.NormalizeDelivery(row, now, row.MaxAttempts)
row, err := notification.NormalizeDelivery(row, now, defaultDeliveryMaxAttempts)
if err != nil {
return notification.DeliveryRow{}, false, err
}
@ -222,7 +225,7 @@ WHERE status = 'leased'
return int(n), nil
}
func (s *Store) MarkDeliverySent(ctx context.Context, id string, externalID string, at time.Time) error {
func (s *Store) MarkDeliverySent(ctx context.Context, id string, owner string, externalID string, at time.Time) error {
if at.IsZero() {
at = time.Now().UTC()
}
@ -234,13 +237,18 @@ SET status = 'sent',
external_id = ?,
delivered_at = ?,
updated_at = ?
WHERE id = ? AND status = 'leased'`, externalID, at, at, id)
WHERE id = ?
AND status = 'leased'
AND lease_owner = ?
AND lease_expires_at > ?`, externalID, at, at, id, owner, at)
}
func (s *Store) MarkDeliveryRetry(ctx context.Context, id string, errCode string, errMessage string, next time.Time) error {
now := time.Now().UTC()
func (s *Store) MarkDeliveryRetry(ctx context.Context, id string, owner string, errCode string, errMessage string, next time.Time, at time.Time) error {
if at.IsZero() {
at = time.Now().UTC()
}
if next.IsZero() {
next = now
next = at
}
return s.updateDelivery(ctx, "mark delivery retry", `UPDATE notification_deliveries
SET attempts = attempts + 1,
@ -251,10 +259,13 @@ SET attempts = attempts + 1,
last_error_code = ?,
last_error = ?,
updated_at = ?
WHERE id = ? AND status = 'leased'`, next, errCode, errMessage, now, id)
WHERE id = ?
AND status = 'leased'
AND lease_owner = ?
AND lease_expires_at > ?`, next, errCode, errMessage, at, id, owner, at)
}
func (s *Store) MarkDeliveryFailed(ctx context.Context, id string, errCode string, errMessage string, at time.Time) error {
func (s *Store) MarkDeliveryFailed(ctx context.Context, id string, owner string, errCode string, errMessage string, at time.Time) error {
if at.IsZero() {
at = time.Now().UTC()
}
@ -266,7 +277,10 @@ SET status = 'failed',
last_error_code = ?,
last_error = ?,
updated_at = ?
WHERE id = ? AND status NOT IN ('sent','failed','skipped','cancelled')`, errCode, errMessage, at, id)
WHERE id = ?
AND status = 'leased'
AND lease_owner = ?
AND lease_expires_at > ?`, errCode, errMessage, at, id, owner, at)
}
func (s *Store) MarkDeliverySkipped(ctx context.Context, id string, reason string, at time.Time) error {
@ -346,9 +360,17 @@ func (s *Store) ListDeliveries(ctx context.Context, filter DeliveryFilter) ([]no
func (s *Store) updateDelivery(ctx context.Context, what string, query string, args ...any) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
if _, err := s.writeDB.ExecContext(ctx, query, args...); err != nil {
res, err := s.writeDB.ExecContext(ctx, query, args...)
if err != nil {
return fmt.Errorf("%s: %w", what, err)
}
affected, err := res.RowsAffected()
if err != nil {
return fmt.Errorf("%s rows affected: %w", what, err)
}
if affected == 0 {
return fmt.Errorf("%s: %w", what, notification.ErrDeliveryUpdateConflict)
}
return nil
}

61
backend/internal/storage/sqlite/notification_delivery_store_test.go
View File

@ -2,6 +2,7 @@ package sqlite
import (
"context"
"errors"
"fmt"
"testing"
"time"
@ -45,6 +46,20 @@ func TestNotificationDeliveryEnqueueIdempotentAndCDC(t *testing.T) {
}
}
func TestNotificationDeliveryEnqueueDefaultMaxAttempts(t *testing.T) {
s, ntf := newDeliveryTestNotification(t, "delivery-default-max")
ctx := context.Background()
row := sampleDelivery(ntf, "desktop")
row.MaxAttempts = 0
got, _, err := s.EnqueueDelivery(ctx, row)
if err != nil {
t.Fatal(err)
}
if got.MaxAttempts != 5 {
t.Fatalf("default max attempts = %d, want 5", got.MaxAttempts)
}
}
func TestNotificationDeliveryClaimDueStableOrder(t *testing.T) {
s, ntf := newDeliveryTestNotification(t, "delivery-claim")
ctx := context.Background()
@ -126,7 +141,7 @@ func TestNotificationDeliveryMarkSentRetryFailedAndSkipped(t *testing.T) {
if len(claimed) != 1 {
t.Fatalf("claim sent row len=%d", len(claimed))
}
if err := s.MarkDeliverySent(ctx, sent.ID, "native-1", now.Add(time.Second)); err != nil {
if err := s.MarkDeliverySent(ctx, sent.ID, "owner", "native-1", now.Add(time.Second)); err != nil {
t.Fatal(err)
}
got, _, _ := s.GetDelivery(ctx, sent.ID)
@ -142,7 +157,7 @@ func TestNotificationDeliveryMarkSentRetryFailedAndSkipped(t *testing.T) {
t.Fatalf("claim retry row len=%d", len(claimed))
}
next := now.Add(30 * time.Second)
if err := s.MarkDeliveryRetry(ctx, retry.ID, "timeout", "timed out", next); err != nil {
if err := s.MarkDeliveryRetry(ctx, retry.ID, "owner", "timeout", "timed out", next, now.Add(time.Second)); err != nil {
t.Fatal(err)
}
got, _, _ = s.GetDelivery(ctx, retry.ID)
@ -158,7 +173,7 @@ func TestNotificationDeliveryMarkSentRetryFailedAndSkipped(t *testing.T) {
if len(claimed) != 1 {
t.Fatalf("claim fail row len=%d", len(claimed))
}
if err := s.MarkDeliveryRetry(ctx, fail.ID, "timeout", "timed out", next); err != nil {
if err := s.MarkDeliveryRetry(ctx, fail.ID, "owner", "timeout", "timed out", next, now.Add(time.Second)); err != nil {
t.Fatal(err)
}
got, _, _ = s.GetDelivery(ctx, fail.ID)
@ -179,8 +194,8 @@ func TestNotificationDeliveryMarkSentRetryFailedAndSkipped(t *testing.T) {
t.Fatalf("skipped row should not be claimable: %+v", claimed)
}
}
if err := s.MarkDeliveryRetry(ctx, skipped.ID, "timeout", "timed out", next); err != nil {
t.Fatal(err)
if err := s.MarkDeliveryRetry(ctx, skipped.ID, "owner", "timeout", "timed out", next, now.Add(time.Second)); !errors.Is(err, notification.ErrDeliveryUpdateConflict) {
t.Fatalf("retry skipped row err = %v, want update conflict", err)
}
got, _, _ = s.GetDelivery(ctx, skipped.ID)
if got.Status != notification.DeliverySkipped || got.Attempts != 0 {
@ -188,6 +203,40 @@ func TestNotificationDeliveryMarkSentRetryFailedAndSkipped(t *testing.T) {
}
}
func TestNotificationDeliveryCompletionFencedByLeaseOwner(t *testing.T) {
s, ntf := newDeliveryTestNotification(t, "delivery-owner-fence")
ctx := context.Background()
now := time.Now().UTC().Truncate(time.Second)
row, _, err := s.EnqueueDelivery(ctx, sampleDueDelivery(ntf, "desktop", now))
if err != nil {
t.Fatal(err)
}
if _, err := s.ClaimDueDeliveries(ctx, notification.SinkAOApp, "owner-1", now, 1, time.Second); err != nil {
t.Fatal(err)
}
if released, err := s.ReleaseExpiredDeliveryLeases(ctx, now.Add(2*time.Second)); err != nil || released != 1 {
t.Fatalf("release = %d err=%v", released, err)
}
if _, err := s.ClaimDueDeliveries(ctx, notification.SinkAOApp, "owner-2", now.Add(2*time.Second), 1, time.Second); err != nil {
t.Fatal(err)
}
if err := s.MarkDeliverySent(ctx, row.ID, "owner-1", "stale", now.Add(2500*time.Millisecond)); !errors.Is(err, notification.ErrDeliveryUpdateConflict) {
t.Fatalf("stale owner MarkDeliverySent err = %v, want update conflict", err)
}
got, _, _ := s.GetDelivery(ctx, row.ID)
if got.Status != notification.DeliveryLeased || got.LeaseOwner != "owner-2" || got.ExternalID != "" {
t.Fatalf("stale owner should not change active lease: %+v", got)
}
if err := s.MarkDeliverySent(ctx, row.ID, "owner-2", "native-2", now.Add(2500*time.Millisecond)); err != nil {
t.Fatalf("current owner sent: %v", err)
}
got, _, _ = s.GetDelivery(ctx, row.ID)
if got.Status != notification.DeliverySent || got.ExternalID != "native-2" {
t.Fatalf("current owner should complete delivery: %+v", got)
}
}
func TestNotificationDeliveryUpdateCDC(t *testing.T) {
s, ntf := newDeliveryTestNotification(t, "delivery-cdc-update")
ctx := context.Background()
@ -199,7 +248,7 @@ func TestNotificationDeliveryUpdateCDC(t *testing.T) {
if _, err := s.ClaimDueDeliveries(ctx, notification.SinkAOApp, "owner", time.Now().UTC(), 1, time.Minute); err != nil {
t.Fatal(err)
}
if err := s.MarkDeliveryFailed(ctx, row.ID, "permanent", "bad route", time.Now().UTC()); err != nil {
if err := s.MarkDeliveryFailed(ctx, row.ID, "owner", "permanent", "bad route", time.Now().UTC()); err != nil {
t.Fatal(err)
}
evs, err := s.ReadChangeLogAfter(ctx, startSeq, 10)