Harden distributed pipeline

coordinator/internal/memstore/memstore.go

@@ -160,17 +160,18 @@ func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time
 	return cancelled, nil
 }
 
-func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
+func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	var n int64
+	affected := make([]uuid.UUID, 0)
 	for _, t := range r.tasks {
-		if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
+		if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
+			t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
 			t.ExpireLease(now)
-			n++
+			affected = append(affected, t.JobID)
 		}
 	}
-	return n, nil
+	return affected, nil
 }
 
 // --- JobRepo -------------------------------------------------------------
@@ -298,6 +299,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
 	return &cp, nil
 }
 
+func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	for _, a := range r.arts {
+		if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult &&
+			a.Attempt != nil && *a.Attempt == attempt {
+			return cloneArtifact(a), nil
+		}
+	}
+	return nil, nil
+}
+
+func cloneArtifact(a *domain.Artifact) *domain.Artifact {
+	cp := *a
+	return &cp
+}
+
 // --- BlobStore -----------------------------------------------------------
 
 type BlobStore struct {

coordinator/internal/memstore/ui_read.go

@@ -60,6 +60,18 @@ func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domai
 	sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
 	return out, nil
 }
+
+func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
+	out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
+	for _, id := range jobIDs {
+		tasks, err := r.ListTasksByJob(ctx, id)
+		if err != nil {
+			return nil, err
+		}
+		out[id] = tasks
+	}
+	return out, nil
+}
 func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
 	if limit < 1 || limit > 100 {
 		return nil, domain.ErrInvalidInput