Finalize bounded typed pipeline implementation

internal/framework/pipeline/runner_concurrent.go

@@ -98,6 +98,7 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 	jobs := make(chan extractJob, workerCount)
 	results := make(chan extractJobResult, workerCount)
 	completions := make(chan laneCompletion, len(states))
+	continuations := make(chan *laneExtractState, workerCount)
 
 	var workers sync.WaitGroup
 	for i := 0; i < workerCount; i++ {
@@ -130,26 +131,46 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 		}
 	}()
 	go func() { workers.Wait(); close(results) }()
+	var continuationWorkers sync.WaitGroup
+	for i := 0; i < workerCount; i++ {
+		continuationWorkers.Add(1)
+		go func() {
+			defer continuationWorkers.Done()
+			for state := range continuations {
+				if err := ctx.Err(); err != nil {
+					completions <- laneCompletion{index: state.index, err: err}
+					continue
+				}
+				laneOutput, err := r.continueLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, state)
+				completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
+			}
+		}()
+	}
 
 	completedOutputs := make([]RunOutput, len(states))
 	var runErrors []orderedRunError
+	var pendingContinuations []*laneExtractState
 	launched, completed := 0, 0
-	launch := func(state *laneExtractState) {
-		launched++
-		go func() {
-			laneOutput, err := r.continueLane(ctx, input, checkpoints, loader, doc, sourceInput, sessionID, chunks, state)
-			completions <- laneCompletion{index: state.index, output: laneOutput, err: err}
-		}()
-	}
 	for _, state := range states {
 		if state.decision.Reused {
-			launch(state)
+			pendingContinuations = append(pendingContinuations, state)
 		}
 	}
 
 	resultChannel := results
-	for resultChannel != nil || completed < launched {
+	for resultChannel != nil || len(pendingContinuations) > 0 || completed < launched {
+		var continuationChannel chan<- *laneExtractState
+		var nextContinuation *laneExtractState
+		if len(pendingContinuations) > 0 && ctx.Err() == nil {
+			continuationChannel = continuations
+			nextContinuation = pendingContinuations[0]
+		} else if ctx.Err() != nil {
+			pendingContinuations = nil
+		}
 		select {
+		case continuationChannel <- nextContinuation:
+			pendingContinuations = pendingContinuations[1:]
+			launched++
 		case result, ok := <-resultChannel:
 			if !ok {
 				resultChannel = nil
@@ -171,7 +192,7 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 					runErrors = append(runErrors, orderedRunError{stage: 0, lane: state.index, chunk: len(chunks), err: err})
 					cancel()
 				} else {
-					launch(state)
+					pendingContinuations = append(pendingContinuations, state)
 				}
 			}
 		case completion := <-completions:
@@ -183,6 +204,8 @@ func (r *Runner) runLanes(parent context.Context, input RunInput, checkpoints Ch
 			}
 		}
 	}
+	close(continuations)
+	continuationWorkers.Wait()
 	for i := range completedOutputs {
 		mergeLaneOutput(&output, completedOutputs[i])
 	}