fix: keep public idempotent triggers atomic

Author: colonelpanic8

internal/services/controllers/task/controller.go

@@ -1071,7 +1071,12 @@ func (tc *TasksControllerImpl) handleProcessInternalEvents(ctx context.Context,
 
 // handleProcessEventTrigger is responsible for inserting tasks into the database based on event triggers.
 func (tc *TasksControllerImpl) handleProcessTaskTrigger(ctx context.Context, tenantId uuid.UUID, payloads [][]byte) error {
-	err := tc.tw.TriggerFromWorkflowNames(ctx, tenantId, msgqueue.JSONConvert[v1.WorkflowNameTriggerOpts](payloads))
+	opts := msgqueue.JSONConvert[v1.WorkflowNameTriggerOpts](payloads)
+	for _, opt := range opts {
+		opt.AllowPartialIdempotency = true
+	}
+
+	err := tc.tw.TriggerFromWorkflowNames(ctx, tenantId, opts)
 
 	if errors.Is(err, v1.ErrIdempotencyKeyAlreadyClaimed) {
 		tc.l.Debug().Err(err).Msg("skipping workflow trigger with duplicate idempotency key")

pkg/repository/idempotency_trigger_test.go

@@ -510,8 +510,8 @@ func TestTriggerWorkflowIdempotency_ReclaimDoesNotUpdateLastDeniedAt(t *testing.
 	assert.False(t, lastDeniedAt.Valid)
 }
 
-func TestTriggerWorkflowIdempotency_MixedKeysAllowsPartial(t *testing.T) {
-	// Mixed batch: one reclaimable, one running. We should allow the reclaimable run and skip the duplicate.
+func TestTriggerWorkflowIdempotency_MixedKeysRejectsPartialByDefault(t *testing.T) {
+	// Public trigger paths should remain atomic: if any key is still duplicate, reject the whole batch.
 	pool, cleanup := setupPostgresWithMigration(t)
 	defer cleanup()
 
@@ -562,6 +562,77 @@ func TestTriggerWorkflowIdempotency_MixedKeysAllowsPartial(t *testing.T) {
 		IdempotencyKey: &keyRunning,
 	}
 
+	_, _, err = repo.Triggers().TriggerFromWorkflowNames(ctx, tenantId, []*WorkflowNameTriggerOpts{terminalSecond, runningSecond})
+	require.Error(t, err)
+	assert.ErrorIs(t, err, ErrIdempotencyKeyAlreadyClaimed)
+
+	claimedTerminal := getClaimedByExternalId(t, pool, tenantId, keyTerminal)
+	require.NotNil(t, claimedTerminal)
+	assert.Equal(t, terminalFirst.ExternalId, *claimedTerminal)
+
+	claimedRunning := getClaimedByExternalId(t, pool, tenantId, keyRunning)
+	require.NotNil(t, claimedRunning)
+	assert.Equal(t, runningFirst.ExternalId, *claimedRunning)
+
+	lastDeniedAt := getLastDeniedAt(t, pool, tenantId, keyRunning)
+	assert.True(t, lastDeniedAt.Valid)
+	assert.WithinDuration(t, time.Now(), lastDeniedAt.Time, 2*time.Minute)
+}
+
+func TestTriggerWorkflowIdempotency_MixedKeysAllowsPartialWhenEnabled(t *testing.T) {
+	// Internal queued trigger processing may skip duplicates while still creating non-duplicate runs.
+	pool, cleanup := setupPostgresWithMigration(t)
+	defer cleanup()
+
+	repo, cleanupRepo := setupRepositoryWithTTL(t, pool, 30*time.Minute, 1*time.Minute)
+	defer cleanupRepo()
+
+	tenantId := setupTenant(t, repo, pool)
+	wf := createMinimalWorkflow(t, pool, tenantId, "test-workflow")
+
+	ctx := context.Background()
+	keyTerminal := IdempotencyKey("idem-key-terminal-allow-partial")
+	keyRunning := IdempotencyKey("idem-key-running-allow-partial")
+
+	terminalFirst := &WorkflowNameTriggerOpts{
+		TriggerTaskData: &TriggerTaskData{
+			WorkflowName: "test-workflow",
+		},
+		ExternalId:     uuid.New(),
+		IdempotencyKey: &keyTerminal,
+	}
+	_, _, err := repo.Triggers().TriggerFromWorkflowNames(ctx, tenantId, []*WorkflowNameTriggerOpts{terminalFirst})
+	require.NoError(t, err)
+	insertRunStatus(t, pool, tenantId, wf, terminalFirst.ExternalId, sqlcv1.V1ReadableStatusOlapFAILED)
+
+	runningFirst := &WorkflowNameTriggerOpts{
+		TriggerTaskData: &TriggerTaskData{
+			WorkflowName: "test-workflow",
+		},
+		ExternalId:     uuid.New(),
+		IdempotencyKey: &keyRunning,
+	}
+	_, _, err = repo.Triggers().TriggerFromWorkflowNames(ctx, tenantId, []*WorkflowNameTriggerOpts{runningFirst})
+	require.NoError(t, err)
+	insertRunStatus(t, pool, tenantId, wf, runningFirst.ExternalId, sqlcv1.V1ReadableStatusOlapRUNNING)
+
+	terminalSecond := &WorkflowNameTriggerOpts{
+		TriggerTaskData: &TriggerTaskData{
+			WorkflowName: "test-workflow",
+		},
+		ExternalId:              uuid.New(),
+		IdempotencyKey:          &keyTerminal,
+		AllowPartialIdempotency: true,
+	}
+	runningSecond := &WorkflowNameTriggerOpts{
+		TriggerTaskData: &TriggerTaskData{
+			WorkflowName: "test-workflow",
+		},
+		ExternalId:              uuid.New(),
+		IdempotencyKey:          &keyRunning,
+		AllowPartialIdempotency: true,
+	}
+
 	_, _, err = repo.Triggers().TriggerFromWorkflowNames(ctx, tenantId, []*WorkflowNameTriggerOpts{terminalSecond, runningSecond})
 	require.NoError(t, err)
 

pkg/repository/ids.go

@@ -30,6 +30,11 @@ type WorkflowNameTriggerOpts struct {
 	// (optional) The idempotency key to use for debouncing this task
 	IdempotencyKey *IdempotencyKey
 
+	// AllowPartialIdempotency permits duplicate keys in the same batch to be skipped while
+	// still creating any non-duplicate runs. Public trigger APIs should keep this disabled
+	// so bulk requests remain atomic.
+	AllowPartialIdempotency bool
+
 	// Whether to skip the creation of the child workflow
 	ShouldSkip bool
 }

pkg/repository/trigger.go

@@ -371,17 +371,7 @@ func (s *sharedRepository) triggerFromWorkflowNames(ctx context.Context, tx *Opt
 	triggerOpts, denyUpdateKeys, duplicateKeys, err := s.prepareTriggerFromWorkflowNames(ctx, tx.tx, tenantId, opts)
 
 	if err != nil {
-		if errors.Is(err, ErrIdempotencyKeyAlreadyClaimed) && len(denyUpdateKeys) > 0 {
-			updateErr := s.queries.UpdateIdempotencyKeysLastDeniedAt(ctx, s.pool, sqlcv1.UpdateIdempotencyKeysLastDeniedAtParams{
-				Tenantid: tenantId,
-				Keys:     denyUpdateKeys,
-			})
-			if updateErr != nil {
-				err = errors.Join(err, fmt.Errorf("failed to update idempotency key deny timestamps: %w", updateErr))
-			}
-		}
-
-		return nil, nil, nil, fmt.Errorf("failed to prepare trigger from workflow names: %w", err)
+		return nil, nil, denyUpdateKeys, fmt.Errorf("failed to prepare trigger from workflow names: %w", err)
 	}
 
 	if len(duplicateKeys) > 0 && len(triggerOpts) > 0 {
@@ -410,6 +400,17 @@ func (r *TriggerRepositoryImpl) TriggerFromWorkflowNames(ctx context.Context, te
 
 	tasks, dags, denyUpdateKeys, err := r.triggerFromWorkflowNames(ctx, tx, tenantId, opts)
 	if err != nil {
+		if errors.Is(err, ErrIdempotencyKeyAlreadyClaimed) && len(denyUpdateKeys) > 0 {
+			tx.Rollback()
+			updateErr := r.queries.UpdateIdempotencyKeysLastDeniedAt(ctx, r.pool, sqlcv1.UpdateIdempotencyKeysLastDeniedAtParams{
+				Tenantid: tenantId,
+				Keys:     denyUpdateKeys,
+			})
+			if updateErr != nil {
+				err = errors.Join(err, fmt.Errorf("failed to update idempotency key deny timestamps: %w", updateErr))
+			}
+		}
+
 		return nil, nil, err
 	}
 
@@ -2325,6 +2326,7 @@ func (r *sharedRepository) prepareTriggerFromWorkflowNames(ctx context.Context,
 
 	// each (workflowVersionId, opt) is a separate workflow that we need to create
 	triggerOpts := make([]triggerTuple, 0, len(opts))
+	allowPartialIdempotency := len(opts) > 0
 
 	for _, workflowVersion := range workflowVersionsByNames {
 		opts, ok := namesToOpts[workflowVersion.WorkflowName]
@@ -2334,6 +2336,8 @@ func (r *sharedRepository) prepareTriggerFromWorkflowNames(ctx context.Context,
 		}
 
 		for _, opt := range opts {
+			allowPartialIdempotency = allowPartialIdempotency && opt.AllowPartialIdempotency
+
 			if opt.IdempotencyKey != nil {
 				keyClaimantPair := KeyClaimantPair{
 					IdempotencyKey:      *opt.IdempotencyKey,
@@ -2367,6 +2371,10 @@ func (r *sharedRepository) prepareTriggerFromWorkflowNames(ctx context.Context,
 		}
 	}
 
+	if len(triggerOpts) > 0 && len(duplicateKeys) > 0 && !allowPartialIdempotency {
+		return nil, denyUpdateKeys, duplicateKeys, &IdempotencyKeyAlreadyClaimedError{Keys: duplicateKeys}
+	}
+
 	if len(triggerOpts) == 0 && len(duplicateKeys) > 0 {
 		return nil, denyUpdateKeys, duplicateKeys, &IdempotencyKeyAlreadyClaimedError{Keys: duplicateKeys}
 	}