I'd like to implement a middleware with the following pseudocode:

// If draft mode was requested, enable drafts and redirect to the same page without draft parameter
if(searchParams.has("draft")) {
  if(searchParams.get("draft") === "true") {
    draftMode().enable();
  } else {
    draftMode().disable();
  }
  const redirectUrl = new URL(request.url);
  redirectUrl.searchParams.delete("draft");
  return NextResponse.redirect(redirectUrl)
}

Unfortunately, at the moment, when I try this, the redirect works, but draft mode is not changed. If I instead redirect user to the /api/draft page which does similar logic, but handles redirect with redirect_url query parameter, the drat changes correctly.
I'd love to keep this logic only in middleware and not use dedicated api routes for this. Is it possible currently? Do I need to pass somehow draftMode to the NextResponse.redirect response? If so, how would I do it?