internal/runtime/executor/dag_runner.go (1)

439-444: ⚠️ Potential issue | 🟠 Major

Track distributed retry runs before waiting.

Kill only requests coordinator cancellation for IDs in distributedRuns. The distributed Execute path records the run before dispatch, but Retry does not, so canceling a parent while a distributed step retry is running can leave the coordinator task uncanceled and wait until timeout.

Proposed fix

 	if e.shouldDispatchToCoordinator(rCtx.DefaultExecMode) {
 		logger.Info(ctx, "Retrying sub DAG via distributed execution", tag.Step(stepName))
+		e.mu.Lock()
+		e.distributedRuns[runParams.RunID] = true
+		e.mu.Unlock()
 		if err := e.dispatchRetryToCoordinator(ctx, runParams, stepName); err != nil {
 			return nil, fmt.Errorf("distributed step retry failed: %w", err)
 		}
 		return e.waitCompletion(ctx, runParams.RunID)
 	}

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@internal/runtime/executor/dag_runner.go` around lines 439 - 444, The retry
path doesn't register the distributed run ID before dispatching, so the Kill
logic (which cancels IDs in distributedRuns) can miss it; modify the retry flow
in the block guarded by e.shouldDispatchToCoordinator to record the run into the
same distributed run tracking used by the Execute path (the distributedRuns
registry/collection on the executor) prior to calling
e.dispatchRetryToCoordinator(ctx, runParams, stepName), and ensure any cleanup
or removal remains paired (e.g., after waitCompletion or on error) so the
coordinator cancellation logic sees the retry run ID.