fix: ux behavior on save for anonymous users

should redirect to login page with next=...

daras_ai_v2/base.py

@@ -95,6 +95,7 @@
 
 
 SUBMIT_AFTER_LOGIN_Q = "submitafterlogin"
+SAVE_AFTER_LOGIN_Q = "saveafterlogin"
 
 
 class RecipeRunState(Enum):
@@ -312,6 +313,11 @@ def refresh_state(self):
             gui.session_state.update(output)
 
     def render(self):
+        if self.should_save_after_login():
+            self.save_or_update_published_run_and_redirect(
+                current_run=self.get_current_sr()
+            )
+
         self.setup_sentry()
 
         if self.get_run_state(gui.session_state) == RecipeRunState.running:
@@ -514,23 +520,35 @@ def _render_published_run_save_buttons(
             )
             publish_modal = gui.Modal("Publish to", key="publish-modal")
             if pressed_save:
-                publish_modal.open()
+                if self.request.user and not self.request.user.is_anonymous:
+                    publish_modal.open()
+                else:
+                    self._save_for_anonymous_user(current_run=current_run)
             if publish_modal.is_open():
                 with publish_modal.container(style={"minWidth": "min(500px, 100vw)"}):
                     self._render_publish_modal(
                         current_run=current_run,
                         published_run=published_run,
                         modal=publish_modal,
-                        is_update_mode=can_edit,
                     )
 
+    def _save_for_anonymous_user(self, current_run: SavedRun):
+        if self._has_request_changed():
+            current_run = self.on_submit()
+            if not current_run:
+                gui.rerun()
+
+        self.ensure_authentication(
+            next_url=current_run.get_app_url(query_params={SAVE_AFTER_LOGIN_Q: "1"}),
+            anon_ok=False,
+        )
+
     def _render_publish_modal(
         self,
         *,
         current_run: SavedRun,
         published_run: PublishedRun,
         modal: gui.Modal,
-        is_update_mode: bool = False,
     ):
         if published_run.is_root() and self.is_current_user_admin():
             with gui.div(className="text-danger"):
@@ -578,11 +596,10 @@ def _render_publish_modal(
                 )
 
         with gui.div(className="mt-4"):
-            if is_update_mode:
+            if self.can_user_edit_published_run(published_run):
                 title = published_run.title or self.title
             else:
-                recipe_title = self.get_root_published_run().title or self.title
-                title = f"{self.request.user.first_name_possesive()} {recipe_title}"
+                title = self._get_default_pr_title()
             published_run_title = gui.text_input(
                 "##### Title",
                 key="published_run_title",
@@ -607,46 +624,75 @@ def _render_publish_modal(
 
         self._render_admin_options(current_run, published_run)
 
-        if not pressed_save:
-            return
+        update_existing_published_run = self.can_user_edit_published_run(published_run)
+        if pressed_save:
+            if not (published_run.is_root() and update_existing_published_run):
+                # don't validate title root example
+                try:
+                    self._validate_published_run_title(
+                        gui.session_state["published_run_title"]
+                    )
+                except TitleValidationError as e:
+                    gui.error(str(e))
+                    return
 
-        is_root_published_run = is_update_mode and published_run.is_root()
-        if not is_root_published_run:
-            try:
-                self._validate_published_run_title(published_run_title)
-            except TitleValidationError as e:
-                gui.error(str(e))
-                return
+            self.save_or_update_published_run_and_redirect(
+                current_run=current_run,
+                published_run=(
+                    published_run if update_existing_published_run else None
+                ),
+                title=published_run_title.strip(),
+                notes=published_run_notes.strip(),
+                visibility=published_run_visibility,
+            )
 
+    def save_or_update_published_run_and_redirect(
+        self,
+        *,
+        current_run: SavedRun,
+        published_run: PublishedRun | None = None,
+        title: str | None = None,
+        notes: str | None = None,
+        visibility: PublishedRunVisibility | None = None,
+    ):
         if self._has_request_changed():
             current_run = self.on_submit()
             if not current_run:
-                modal.close()
+                gui.rerun()
 
-        if is_update_mode:
+        if published_run:
+            # update existing published run
             updates = dict(
                 saved_run=current_run,
-                title=published_run_title.strip(),
-                notes=published_run_notes.strip(),
-                visibility=published_run_visibility,
+                title=title or published_run.title,
+                notes=(
+                    notes if notes is not None else published_run.notes
+                ),  # allow notes = ""
+                visibility=visibility or published_run.visibility,
             )
-            if not self._has_published_run_changed(
-                published_run=published_run, **updates
-            ):
+            if not self._has_published_run_changed(published_run, **updates):
                 gui.error("No changes to publish", icon="⚠️")
                 return
             published_run.add_version(user=self.request.user, **updates)
         else:
+            # save new published run
             published_run = self.create_published_run(
                 published_run_id=get_random_doc_id(),
                 saved_run=current_run,
                 user=self.request.user,
-                title=published_run_title.strip(),
-                notes=published_run_notes.strip(),
-                visibility=published_run_visibility,
+                title=title or self._get_default_pr_title(),
+                notes=notes or "",
+                visibility=visibility or PublishedRunVisibility.UNLISTED,
             )
         raise gui.RedirectException(published_run.get_app_url())
 
+    def _get_default_pr_title(self) -> str:
+        recipe_title = self.get_root_published_run().title or self.title
+        if self.request.user:
+            return f"{self.request.user.first_name_possesive()} {recipe_title}"
+        else:
+            return f"My {recipe_title}"
+
     def _validate_published_run_title(self, title: str):
         if slugify(title) in settings.DISALLOWED_TITLE_SLUGS:
             raise TitleValidationError(
@@ -661,8 +707,8 @@ def _validate_published_run_title(self, title: str):
 
     def _has_published_run_changed(
         self,
-        *,
         published_run: PublishedRun,
+        *,
         saved_run: SavedRun,
         title: str,
         notes: str,
@@ -1619,6 +1665,14 @@ def should_submit_after_login(self) -> bool:
             and not self.request.user.is_anonymous
         )
 
+    def should_save_after_login(self) -> bool:
+        return (
+            gui.get_query_params().get(SAVE_AFTER_LOGIN_Q)
+            and self.request
+            and self.request.user
+            and not self.request.user.is_anonymous
+        )
+
     def create_new_run(
         self, *, enable_rate_limits: bool = False, **defaults
     ) -> SavedRun: