with-task-group.swift

// A fix for https://github.com/AvdLee/TaskGroupsResultBuilder/blob/main/TaskGroups/TaskResultBuilder.swift
// to handle cancellation.

@resultBuilder
struct TaskBuilder {
  static func buildExpression<Success: Sendable>(_ task: Task<Success, Never>) -> [Task<Success, Never>] {
    [task]
  }

  static func buildBlock<Success: Sendable>(_ tasks: [Task<Success, Never>]...) -> [Task<Success, Never>] {
    tasks.flatMap { $0 }
  }

  static func buildFinalResult<Success: Sendable>(_ tasks: [Task<Success, Never>]) -> Task<[Success], Never> {
    Task {
      await withTaskGroup(of: Success.self, returning: [Success].self) { taskGroup in
        tasks.forEach { task in
          taskGroup.addTask {
            await withTaskCancellationHandler { // 👈
              await task.value
            } onCancel: {
              task.cancel()
            }
          }
        }

        return await taskGroup.reduce(into: [Success]()) { partialResult, name in
          partialResult.append(name)
        }
      }
    }
  }
}

func withTaskGroup<Success: Sendable>(@TaskBuilder builder: () -> Task<[Success], Never>) async -> [Success] {
  let task = builder()
  return await withTaskCancellationHandler { // 👈
    await task.value
  } onCancel: {
    task.cancel()
  }
}

// Test to confirm:

func testTaskGroupCancellation() async throws {
  let task = Task {
    await withTaskGroup {
      Task {
        do {
          try await Task.sleep(for: .seconds(1))
          XCTFail("🛑 This should not be executed")
        } catch {}
      }
    }
  }
  try await Task.sleep(for: .milliseconds(100))
  task.cancel()
  try await Task.sleep(for: .seconds(1))
}